add invite system

app/routes/users.py

@@ -4,7 +4,7 @@ from flask import (
 from functools import wraps
 from ..db import db
 from ..lib.babycode import babycode_to_html
-from ..models import Users, Sessions, Subscriptions, Avatars, PasswordResetLinks
+from ..models import Users, Sessions, Subscriptions, Avatars, PasswordResetLinks, InviteKeys
 from ..constants import InfoboxKind, PermissionLevel
 from ..auth import digest, verify
 from wand.image import Image
@@ -195,32 +195,53 @@ def log_in_post():
 @bp.get("/sign_up")
 @redirect_if_logged_in(".page", username = lambda: get_active_user().username)
 def sign_up():
+    if current_app.config['DISABLE_SIGNUP']:
+        key = request.args.get('key', default=None)
+        if key is None:
+            return redirect(url_for('topics.all_topics'))
+
+        invite = InviteKeys.find({'key': key})
+        if not invite:
+            return redirect(url_for('topics.all_topics'))
+
+        inviter = Users.find({'id': invite.created_by})
+        return render_template("users/sign_up.html", inviter=inviter, key=key)
+
     return render_template("users/sign_up.html")
 
 
 @bp.post("/sign_up")
 @redirect_if_logged_in(".page", username = lambda: get_active_user().username)
 def sign_up_post():
+    key = request.form.get('key', default=None)
+
+    if current_app.config['DISABLE_SIGNUP']:
+        if not key:
+            return redirect(url_for("topics.all_topics"))
+        invite_key = InviteKeys.find({'key': key})
+        if not invite_key:
+            return redirect(url_for("topics.all_topics"))
+
     username = request.form['username']
     password = request.form['password']
     password_confirm = request.form['password-confirm']
 
     if not validate_username(username):
         flash("Invalid username.", InfoboxKind.ERROR)
-        return redirect(url_for("users.sign_up"))
+        return redirect(url_for("users.sign_up", key=key))
 
     user_exists = Users.count({"username": username}) > 0
     if user_exists:
         flash(f"Username '{username}' is already taken.", InfoboxKind.ERROR)
-        return redirect(url_for("users.sign_up"))
+        return redirect(url_for("users.sign_up", key=key))
 
     if not validate_password(password):
         flash("Invalid password.", InfoboxKind.ERROR)
-        return redirect(url_for("users.sign_up"))
+        return redirect(url_for("users.sign_up", key=key))
 
     if password != password_confirm:
         flash("Passwords do not match.", InfoboxKind.ERROR)
-        return redirect(url_for("users.sign_up"))
+        return redirect(url_for("users.sign_up", key=key))
 
     hashed = digest(password)
 
@@ -230,11 +251,19 @@ def sign_up_post():
         "permission": PermissionLevel.GUEST.value,
     })
 
+    if current_app.config['DISABLE_SIGNUP']:
+        invite_key = InviteKeys.find({'key': key})
+        new_user.update({
+            'invited_by': invite_key.created_by,
+            'permission': PermissionLevel.USER.value,
+        })
+        invite_key.delete()
+
     session_obj = create_session(new_user.id)
 
     session['pyrom_session_key'] = session_obj.key
     flash("Signed up successfully!", InfoboxKind.INFO)
-    return redirect(url_for("users.sign_up"))
+    return redirect(url_for("topics.all_topics"))
 
 
 @bp.get("/<username>")
@@ -573,3 +602,69 @@ def reset_link_login_form(key):
     flash("Logged in!", InfoboxKind.INFO)
 
     return redirect(url_for('.page', username=target_user.username))
+
+
+@bp.get('/<username>/invite-links/')
+@login_required
+def invite_links(username):
+    target_user = Users.find({
+        'username': username
+    })
+    if not target_user or not target_user.can_invite():
+        return redirect(url_for('.page', username=username))
+
+    if target_user.username != get_active_user().username:
+        return redirect(url_for('.invite_links', username=target_user.username))
+
+    invites = InviteKeys.findall({
+        'created_by': target_user.id
+    })
+
+    return render_template('users/invite_links.html', invites=invites)
+
+
+@bp.post('/<username>/invite-links/create')
+@login_required
+def create_invite_link(username):
+    target_user = Users.find({
+        'username': username
+    })
+    if not target_user or not target_user.can_invite():
+        return redirect(url_for('.page', username=username))
+
+    if target_user.username != get_active_user().username:
+        return redirect(url_for('.invite_links', username=target_user.username))
+
+    invite = InviteKeys.create({
+        'created_by': target_user.id,
+        'key': secrets.token_urlsafe(20),
+    })
+
+    return redirect(url_for('.invite_links', username=target_user.username))
+
+
+@bp.post('/<username>/invite-links/revoke')
+@login_required
+def revoke_invite_link(username):
+    target_user = Users.find({
+        'username': username
+    })
+    if not target_user or not target_user.can_invite():
+        return redirect(url_for('.page', username=username))
+
+    if target_user.username != get_active_user().username:
+        return redirect(url_for('.invite_links', username=target_user.username))
+
+    invite = InviteKeys.find({
+        'key': request.form.get('key'),
+    })
+
+    if not invite:
+        return redirect(url_for('.invite_links', username=target_user.username))
+
+    if invite.created_by != target_user.id:
+        return redirect(url_for('.invite_links', username=target_user.username))
+
+    invite.delete()
+
+    return redirect(url_for('.invite_links', username=target_user.username))