SSH: account creation and key upload

dojo_plugin/__init__.py

@@ -30,6 +30,7 @@
 from .pages.sensai import sensai
 from .pages.users import users
 from .pages.settings import settings_override
+from .pages.ssh_key import ssh_key
 from .pages.discord import discord
 from .pages.course import course
 from .pages.belts import belts
@@ -190,6 +191,7 @@ def publish_stat_events_after_request(response):
     app.register_blueprint(belts)
     app.register_blueprint(research)
     app.register_blueprint(feed)
+    app.register_blueprint(ssh_key)
     app.register_blueprint(test_error_pages)
     app.register_blueprint(api, url_prefix="/pwncollege_api/v1")
 

dojo_plugin/api/v1/auth.py

@@ -1,4 +1,5 @@
 import datetime
+import secrets
 from flask import request, session
 from flask_restx import Namespace, Resource
 from CTFd.models import Users, UserFieldEntries, UserFields, db
@@ -9,8 +10,13 @@
 from CTFd.utils.config import can_send_mail
 from CTFd.utils.security.signing import unserialize
 from itsdangerous.exc import BadSignature, BadTimeSignature, SignatureExpired
+from sqlalchemy.exc import IntegrityError
 import base64
 
+from .user import is_ssh_service_request, ssh_service_token
+from ...models import SSHKeys
+from ...utils.ssh_key import InvalidKeyError, normalize_offered_ssh_key
+
 auth_namespace = Namespace("auth", description="Authentication endpoints")
 
 
@@ -28,8 +34,14 @@ def post(self):
         if not registration_visible():
             return {"success": False, "errors": ["Registration is currently disabled"]}, 403
 
-        req = request.get_json()
+        req = request.get_json() or {}
         errors = []
+        ssh_service_request = is_ssh_service_request()
+        if ssh_service_request:
+            try:
+                assert ssh_service_token() == "ssh-onboarding"
+            except Exception:
+                return {"success": False, "error": "Failed to authenticate SSH service token."}, 401
 
         # Get registration data
         name = req.get("name", "").strip()
@@ -39,6 +51,13 @@ def post(self):
         affiliation = req.get("affiliation")
         country = req.get("country")
 
+        ssh_key = None
+        if ssh_service_request:
+            try:
+                ssh_key = normalize_offered_ssh_key(req.get("key_type", ""), req.get("key_base64", ""))
+            except (InvalidKeyError, NotImplementedError) as error:
+                errors.append(f"Invalid SSH key: {error}")
+
         # Check user limit
         num_users_limit = int(get_config("num_users", default=0))
         num_users = Users.query.filter_by(banned=False, hidden=False).count()
@@ -60,10 +79,14 @@ def post(self):
         if not email.check_email_is_whitelisted(email_address):
             errors.append("Email address is not from an allowed domain")
 
-        if len(password) == 0:
-            errors.append("Please provide a password")
-        if len(password) > 128:
-            errors.append("Password is too long")
+        if ssh_service_request:
+            if get_config("registration_code"):
+                errors.append("SSH registration is not available while a registration code is required")
+        else:
+            if len(password) == 0:
+                errors.append("Please provide a password")
+            if len(password) > 128:
+                errors.append("Password is too long")
 
         if website and not validators.validate_url(website):
             errors.append("Website must be a valid URL")
@@ -78,7 +101,7 @@ def post(self):
             errors.append("Affiliation is too long")
 
         # Check registration code if required
-        if get_config("registration_code"):
+        if not ssh_service_request and get_config("registration_code"):
             registration_code = req.get("registration_code", "")
             if registration_code.lower() != str(get_config("registration_code", "")).lower():
                 errors.append("Invalid registration code")
@@ -95,26 +118,29 @@ def post(self):
             return {"success": False, "errors": errors}, 400
 
         # Create user
-        user = Users(name=name, email=email_address, password=password)
+        user = Users(
+            name=name,
+            email=email_address,
+            password=secrets.token_urlsafe(48) if ssh_service_request else password,
+        )
         if website:
             user.website = website
         if affiliation:
             user.affiliation = affiliation
         if country:
             user.country = country
 
-        db.session.add(user)
-        db.session.commit()
-
-        # Add custom field entries
-        for field_id, value in fields.items():
-            entry = UserFieldEntries(
-                field_id=field_id,
-                value=value,
-                user_id=user.id
-            )
-            db.session.add(entry)
-        db.session.commit()
+        try:
+            db.session.add(user)
+            db.session.flush()
+            if ssh_key:
+                db.session.add(SSHKeys(user_id=user.id, value=ssh_key))
+            for field_id, value in fields.items():
+                db.session.add(UserFieldEntries(field_id=field_id, value=value, user_id=user.id))
+            db.session.commit()
+        except IntegrityError:
+            db.session.rollback()
+            return {"success": False, "errors": ["SSH key, username, or email is already in use"]}, 400
 
         # Send verification email if configured
         if get_config("verify_emails") and can_send_mail():
@@ -127,7 +153,18 @@ def post(self):
             if can_send_mail():
                 email.successful_registration_notification(user.email)
 
-        # Set session
+        if ssh_service_request:
+            session.clear()
+            return {
+                "success": True,
+                "user": {
+                    "id": user.id,
+                    "name": user.name,
+                    "email": user.email,
+                    "verified": verified
+                }
+            }
+
         session["id"] = user.id
         session["name"] = user.name
         session["type"] = user.type
@@ -324,4 +361,4 @@ def post(self, token):
         return {
             "success": True,
             "data": {"message": "Password successfully reset"}
-        }
+        }

dojo_plugin/api/v1/ssh_key.py

@@ -1,43 +1,52 @@
-import re
+import datetime
+import secrets
 
 from flask import request
 from flask_restx import Namespace, Resource
 from sqlalchemy.exc import IntegrityError
 from CTFd.models import db
 from CTFd.utils.decorators import authed_only
 from CTFd.utils.user import get_current_user
-from sshpubkeys import SSHKey, InvalidKeyError
-import base64
 import markupsafe
 
-from ...models import SSHKeys
+from ...config import DOJO_HOST
+from ...models import SSHKeyLinkRequests, SSHKeys
+from ...utils.ssh_key import InvalidKeyError, normalize_offered_ssh_key, normalize_ssh_key
+from ...utils.ssh_onboarding import ssh_link_token_digest
+from .user import ssh_service_only
 
 
 ssh_key_namespace = Namespace(
     "keys", description="Endpoint to manage users' public SSH keys"
 )
+LINK_TTL = datetime.timedelta(minutes=15)
+
+
+def link_url(token):
+    host = DOJO_HOST or request.host
+    scheme = "http" if host.startswith("localhost") or host.startswith("127.") else "https"
+    return f"{scheme}://{host}/ssh/link/{token}"
 
 
 @ssh_key_namespace.route("")
 class UpdateKey(Resource):
     @authed_only
     def post(self):
-        data = request.get_json()
-        key_value = data.get("ssh_key", "")
-
-        if key_value:
-            try:
-                key = SSHKey(key_value, strict=True)
-                key.parse()
-                key_value = f"{key.key_type.decode()} {base64.b64encode(key._decoded_key).decode()}"
-            except (InvalidKeyError, NotImplementedError) as e:
-                return (
-                    {
-                        "success": False,
-                        "error": f"Invalid SSH Key, error: <code>{markupsafe.escape(e)}</code> <br>Refer below for how to generate a valid ssh key"
-                    },
-                    400,
-                )
+        data = request.get_json() or {}
+        key_value = data.get("ssh_key", "").strip()
+        if not key_value:
+            return {"success": False, "error": "Please provide an SSH key"}, 400
+
+        try:
+            key_value = normalize_ssh_key(key_value)
+        except (InvalidKeyError, NotImplementedError) as e:
+            return (
+                {
+                    "success": False,
+                    "error": f"Invalid SSH Key, error: <code>{markupsafe.escape(e)}</code> <br>Refer below for how to generate a valid ssh key"
+                },
+                400,
+            )
 
         user = get_current_user()
 
@@ -56,12 +65,19 @@ def post(self):
 
     @authed_only
     def delete(self):
-        data = request.get_json()
-        key_value = data.get("ssh_key", "")
+        data = request.get_json() or {}
+        key_value = data.get("ssh_key", "").strip()
+        if not key_value:
+            return {"success": False, "error": "Please provide an SSH key"}, 400
+
+        try:
+            key_value = normalize_ssh_key(key_value)
+        except (InvalidKeyError, NotImplementedError):
+            return {"success": False, "error": "SSH Key does not exist"}, 400
 
         user = get_current_user()
 
-        key = SSHKeys.query.filter_by(user=user, value=key_value).first()
+        key = SSHKeys.query.filter_by(user_id=user.id, value=key_value).first()
         if not key:
             return (
                 {"success": False, "error": "SSH Key does not exist"},
@@ -72,3 +88,54 @@ def delete(self):
         db.session.commit()
 
         return {"success": True}
+
+
+@ssh_key_namespace.route("/link")
+class CreateLink(Resource):
+    @ssh_service_only
+    def post(self):
+        data = request.get_json() or {}
+        try:
+            key_value = normalize_offered_ssh_key(data.get("key_type", ""), data.get("key_base64", ""))
+        except (InvalidKeyError, NotImplementedError) as error:
+            return {"success": False, "error": f"Invalid SSH key: {error}"}, 400
+
+        if SSHKeys.query.filter_by(value=key_value).first():
+            return {"success": False, "error": "SSH key is already linked to an account"}, 400
+
+        token = secrets.token_urlsafe(32)
+        link_request = SSHKeyLinkRequests(
+            token_digest=ssh_link_token_digest(token),
+            key_value=key_value,
+            fingerprint=data.get("fingerprint", ""),
+            expiration=datetime.datetime.utcnow() + LINK_TTL,
+        )
+        db.session.add(link_request)
+        db.session.commit()
+        return {
+            "success": True,
+            "token": token,
+            "link_url": link_url(token),
+            "expires_at": link_request.expiration.isoformat(),
+        }
+
+
+@ssh_key_namespace.route("/link/<token>")
+class LinkStatus(Resource):
+    @ssh_service_only
+    def get(self, token):
+        link_request = SSHKeyLinkRequests.query.filter_by(token_digest=ssh_link_token_digest(token)).first()
+        if not link_request:
+            return {"success": False, "status": "not_found"}, 404
+        if link_request.consumed:
+            return {
+                "success": True,
+                "status": "linked",
+                "user": {
+                    "id": link_request.user.id if link_request.user else None,
+                    "name": link_request.user.name if link_request.user else None,
+                },
+            }
+        if link_request.expiration < datetime.datetime.utcnow():
+            return {"success": True, "status": "expired"}
+        return {"success": True, "status": "pending"}

dojo_plugin/api/v1/user.py

@@ -1,29 +1,59 @@
 from flask_restx import Namespace, Resource
-from flask import current_app, request, session
+from flask import current_app, g, request, session
 from itsdangerous.url_safe import URLSafeTimedSerializer
 from CTFd.utils.decorators import authed_only
 from CTFd.utils.user import get_current_user
 from CTFd.models import Users
+from functools import wraps
 from ...config import DOJO_SSH_SERVICE_KEY
 from ...utils import get_current_container
 
 user_namespace = Namespace("user", description="User management endpoints")
 CLI_AUTH_PREFIX = "sk-workspace-local-"
 SSH_AUTH_PREFIX = "sk-ssh-service-"
 
-def authed_only_ssh(func):
+def ssh_service_token():
+    auth_header = request.headers.get("Authorization", "")
+    if not auth_header.startswith("Bearer "):
+        return None
+    token = auth_header[len("Bearer "):].strip()
+    if not token.startswith(SSH_AUTH_PREFIX):
+        return None
+    token = token[len(SSH_AUTH_PREFIX):].strip()
+    return URLSafeTimedSerializer(DOJO_SSH_SERVICE_KEY).loads(token, max_age=300)
+
+
+def is_ssh_service_request():
+    auth_header = request.headers.get("Authorization", "")
+    return auth_header.startswith(f"Bearer {SSH_AUTH_PREFIX}")
+
+
+def ssh_service_only(func):
+    @wraps(func)
     def wrapper(*args, **kwargs):
-        auth_header = request.headers.get("Authorization")
-        if not auth_header:
-            return func(*args, **kwargs)
-        if not auth_header.startswith("Bearer "):
+        if not is_ssh_service_request():
+            return {"success": False, "error": "Missing SSH service token."}, 401
+        try:
+            token_tag = ssh_service_token()
+            assert token_tag == "ssh-onboarding"
+        except Exception:
+            return {"success": False, "error": "Failed to authenticate SSH service token."}, 401
+        g.ssh_service = True
+        try:
             return func(*args, **kwargs)
-        token = auth_header[len("Bearer "):].strip()
-        if not token.startswith(SSH_AUTH_PREFIX):
+        finally:
+            if hasattr(g, "ssh_service"):
+                del g.ssh_service
+    return wrapper
+
+
+def authed_only_ssh(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if not is_ssh_service_request():
             return func(*args, **kwargs)
-        token = token[len(SSH_AUTH_PREFIX):].strip()
         try:
-            user_id, token_tag = URLSafeTimedSerializer(DOJO_SSH_SERVICE_KEY).loads(token, max_age=300)
+            user_id, token_tag = ssh_service_token()
             assert token_tag == "ssh-tui"
         except Exception:
             return {"success": False, "error": "Failed to authenticate ssh service token."}, 401
@@ -39,9 +69,8 @@ def wrapper(*args, **kwargs):
             })
             return func(*args, **kwargs)
         finally:
-            for k in ("id", "name", "type", "verified"):
-                session.pop(k, None)
-        return func(*args, **kwargs)
+            for key in ("id", "name", "type", "verified"):
+                session.pop(key, None)
     return wrapper