Initial commitHEADmaster

9 files changed, 537 insertions, 0 deletions

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..0809d3d
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,11 @@
+from .models import db
+from .oidc import oidc_blueprint
+from .admin import oidc_admin_blueprint
+
+def load(app):
+    db.init_app(app)
+    with app.app_context():
+        db.create_all()
+
+    app.register_blueprint(oidc_blueprint)
+    app.register_blueprint(oidc_admin_blueprint)
diff --git a/admin.py b/admin.py
new file mode 100644
index 0000000..699958b
--- /dev/null
+++ b/admin.py
@@ -0,0 +1,62 @@
+import secrets
+from flask import Blueprint, render_template, request, redirect, url_for
+from CTFd.utils.decorators import admins_only
+from .models import db, OIDCClient, OIDCKey
+from .crypto import generate_rsa_key
+from .config import get_config, set_config
+
+oidc_admin_blueprint = Blueprint(
+    "oidc_admin",
+    __name__,
+    url_prefix="/admin/oidc",
+    template_folder="templates",
+)
+
+
+@oidc_admin_blueprint.route("/", methods=["GET"])
+@admins_only
+def index():
+    clients = OIDCClient.query.all()
+    keys = OIDCKey.query.order_by(OIDCKey.created.desc()).all()
+    baseUrl = get_config("base_url", "")
+    return render_template("index.html", clients=clients, keys=keys, baseUrl=baseUrl)
+
+
+@oidc_admin_blueprint.route("/config", methods=["POST"])
+@admins_only
+def config():
+    set_config("base_url", request.form["base_url"])
+    return redirect(url_for("oidc_admin.index"))
+
+
+@oidc_admin_blueprint.route("/clients", methods=["POST"])
+@admins_only
+def clients():
+    client = OIDCClient(
+        client_id=request.form["client_id"],
+        client_secret=secrets.token_urlsafe(32),
+        redirect_uris=request.form["redirect_uris"],
+    )
+    db.session.add(client)
+
+    private_pem, public_pem = generate_rsa_key()
+
+    oidc_key = OIDCKey(
+        kid=secrets.token_urlsafe(16),
+        private_pem=private_pem.decode(),
+        public_pem=public_pem.decode(),
+        client_id=client.client_id,
+    )
+    db.session.add(oidc_key)
+    db.session.commit()
+
+    return redirect(url_for("oidc_admin.index"))
+
+
+@oidc_admin_blueprint.route("/clients/<client_id>/delete", methods=["POST"])
+@admins_only
+def delete_client(client_id):
+    client = OIDCClient.query.get_or_404(client_id)
+    db.session.delete(client)
+    db.session.commit()
+    return redirect(url_for("oidc_admin.index"))
diff --git a/config.json b/config.json
new file mode 100644
index 0000000..3d71ee1
--- /dev/null
+++ b/config.json
@@ -0,0 +1,4 @@
+{
+    "name": "OIDC IdP",
+    "route": "/admin/oidc/"
+}
diff --git a/config.py b/config.py
new file mode 100644
index 0000000..ed7d541
--- /dev/null
+++ b/config.py
@@ -0,0 +1,14 @@
+from .models import db, OIDCConfig
+
+def set_config(key, value):
+    cfg = OIDCConfig.query.filter_by(key=key).first()
+    if cfg:
+        cfg.value = value
+    else:
+        cfg = OIDCConfig(key=key, value=value)
+        db.session.add(cfg)
+    db.session.commit()
+
+def get_config(key, default=None):
+    cfg = OIDCConfig.query.filter_by(key=key).first()
+    return cfg.value if cfg else default
diff --git a/crypto.py b/crypto.py
new file mode 100644
index 0000000..a278c25
--- /dev/null
+++ b/crypto.py
@@ -0,0 +1,45 @@
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives import serialization
+from .models import db, OIDCKey, OIDCClient
+import jwt
+import datetime
+
+
+def generate_rsa_key():
+    key = rsa.generate_private_key(
+        public_exponent=65537,
+        key_size=2048,
+    )
+
+    private_pem = key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    public_pem = key.public_key().public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+
+    return private_pem, public_pem
+
+
+def sign_jwt(payload, issuer, client_id):
+    key = OIDCKey.query.filter_by(client_id=client_id).order_by(OIDCKey.created.desc()).first()
+    if not key:
+        raise ValueError("No RSA key is active for this client")
+
+    claims = {
+        "sub": payload["sub"],
+        "iss": issuer,
+        "aud": [client_id],
+        "iat": datetime.datetime.utcnow(),
+        "exp": datetime.datetime.utcnow() + datetime.timedelta(hours=1),
+        "email": payload.get("email", ""),
+        "name": payload.get("name", ""),
+    }
+
+    token = jwt.encode(claims, key.private_pem.encode("utf-8"), algorithm="RS256")
+
+    return token
diff --git a/models.py b/models.py
new file mode 100644
index 0000000..9fcbccd
--- /dev/null
+++ b/models.py
@@ -0,0 +1,79 @@
+from flask_sqlalchemy import SQLAlchemy
+from CTFd.models import Users
+import time
+import secrets
+
+db = SQLAlchemy()
+
+
+class OIDCConfig(db.Model):
+    __tablename__ = "oidc_config"
+
+    id = db.Column(db.Integer, primary_key=True)
+    key = db.Column(db.String(64), unique=True, nullable=False)
+    value = db.Column(db.Text, nullable=False)
+    created = db.Column(db.Integer, default=lambda: int(time.time()))
+
+
+class OIDCClient(db.Model):
+    __tablename__ = 'oidc_client'
+
+    client_id = db.Column(db.String(32), primary_key=True)
+    client_secret = db.Column(db.String(128), nullable=True)
+    redirect_uris = db.Column(db.Text, nullable=False)
+    pkce = db.Column(db.Boolean, default=False, nullable=False)
+    created = db.Column(db.Integer, default=lambda: int(time.time()))
+
+
+class OIDCAuthCode(db.Model):
+    __tablename__ = 'oidc_auth_code'
+
+    id = db.Column(db.Integer, primary_key=True)
+    code = db.Column(db.String(128), unique=True, nullable=False)
+    user_id = db.Column(db.Integer, db.ForeignKey(Users.id, ondelete='CASCADE'), nullable=False)
+    client_id = db.Column(db.String(32), db.ForeignKey(OIDCClient.client_id, ondelete='CASCADE'), nullable=False)
+    redirect_uri = db.Column(db.Text, nullable=False)
+    code_challenge = db.Column(db.String(128), nullable=True)
+    exp = db.Column(db.Integer, nullable=False)
+
+    user = db.relationship(Users, backref=db.backref("oidc_auth_code", uselist=False, lazy="select"))
+    client = db.relationship(OIDCClient, backref=db.backref("oidc_auth_code", uselist=False, lazy="select"))
+
+
+class OIDCRefreshToken(db.Model):
+    __tablename__ = 'oidc_refresh_token'
+
+    id = db.Column(db.Integer, primary_key=True)
+    refresh_token = db.Column(db.String(128), unique=True, nullable=False)
+    user_id = db.Column(db.Integer, db.ForeignKey(Users.id, ondelete='CASCADE'), nullable=False)
+    client_id = db.Column(db.String(32), db.ForeignKey(OIDCClient.client_id, ondelete='CASCADE'), nullable=False)
+    exp = db.Column(db.Integer, nullable=False)
+
+    user = db.relationship(Users, backref=db.backref("oidc_refresh_token", uselist=False, lazy="select"))
+    client = db.relationship(OIDCClient, backref=db.backref("oidc_refresh_token", uselist=False, lazy="select"))
+
+
+class OIDCAccessToken(db.Model):
+    __tablename__ = 'oidc_access_token'
+
+    id = db.Column(db.Integer, primary_key=True)
+    access_token = db.Column(db.String(128), unique=True, nullable=False)
+    user_id = db.Column(db.Integer, db.ForeignKey(Users.id, ondelete='CASCADE'), nullable=False)
+    client_id = db.Column(db.String(32), db.ForeignKey(OIDCClient.client_id, ondelete='CASCADE'), nullable=False)
+    exp = db.Column(db.Integer, nullable=False)
+
+    user = db.relationship(Users, backref=db.backref("oidc_access_token", uselist=False, lazy="select"))
+    client = db.relationship(OIDCClient, backref=db.backref("oidc_access_token", uselist=False, lazy="select"))
+
+
+class OIDCKey(db.Model):
+    __tablename__ = 'oidc_key'
+
+    id = db.Column(db.Integer, primary_key=True)
+    kid = db.Column(db.String(32), unique=True, nullable=False)
+    private_pem = db.Column(db.Text, nullable=False)
+    public_pem = db.Column(db.Text, nullable=False)
+    created = db.Column(db.Integer, default=lambda: int(time.time()))
+    client_id = db.Column(db.String(32), db.ForeignKey(OIDCClient.client_id, ondelete='CASCADE'), nullable=False)
+
+    client = db.relationship(OIDCClient, backref=db.backref("oidc_key", uselist=False, lazy="select"))
diff --git a/oidc.py b/oidc.py
new file mode 100644
index 0000000..71a187c
--- /dev/null
+++ b/oidc.py
@@ -0,0 +1,206 @@
+import time
+import secrets
+import hashlib
+import base64
+import jwt
+import json
+import urllib.parse
+from flask import Blueprint, request, redirect, jsonify
+from CTFd.utils.user import authed, get_current_user
+from CTFd.models import Users
+from CTFd.plugins import bypass_csrf_protection
+from cryptography.hazmat.primitives import serialization
+from .models import db, OIDCClient, OIDCAuthCode, OIDCRefreshToken, OIDCAccessToken, OIDCKey
+from .crypto import sign_jwt
+from .config import get_config
+
+oidc_blueprint = Blueprint(
+    "oidc",
+    __name__,
+    url_prefix="/oidc"
+)
+
+
+@oidc_blueprint.route("/.well-known/openid-configuration")
+def discovery():
+    issuer = get_config("base_url", "")
+    return jsonify({
+        "issuer": f"{issuer}/oidc/",
+        "authorization_endpoint": f"{issuer}/oidc/authorize",
+        "token_endpoint": f"{issuer}/oidc/token",
+        "userinfo_endpoint": f"{issuer}/oidc/userinfo",
+        "jwks_uri": f"{issuer}/oidc/jwks",
+        "response_types_supported": ["code"],
+        "grant_types_supported": ["authorization_code", "refresh_token"],
+        "code_challenge_methods_supported": ["S256"],
+        "token_endpoint_auth_methods_supported": ["client_secret_basic"],
+        "id_token_signing_alg_values_supported": ["RS256"],
+        "scopes_supported": ["openid", "profile", "email"], #TODO maybe support these for real
+    })
+
+
+@oidc_blueprint.route("/jwks")
+def jwks():
+    keys = []
+
+    for key in OIDCKey.query.all():
+        public_key = serialization.load_pem_public_key(key.public_pem.encode("utf-8"))
+        jwk = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(public_key))
+        jwk["kid"] = key.kid
+        keys.append(jwk)
+
+    return jsonify({"keys": keys})
+
+
+@oidc_blueprint.route("/authorize")
+def authorize():
+    if not authed():
+        return redirect(f"/login?next={urllib.parse.quote(request.full_path)}")
+
+    client_id = request.args.get("client_id")
+    redirect_uri = request.args.get("redirect_uri")
+    state = request.args.get("state")
+    code_challenge = request.args.get("code_challenge")
+
+    client = OIDCClient.query.filter_by(client_id=client_id).first()
+    if not client or redirect_uri not in str.splitlines(client.redirect_uris):
+        return "Invalid client", 400
+
+    if not code_challenge and client.pkce:
+        return "PKCE required", 400
+
+    code = secrets.token_urlsafe(32)
+
+    db.session.add(OIDCAuthCode(
+        code=code,
+        user_id=get_current_user().id,
+        client_id=client.client_id,
+        redirect_uri=redirect_uri,
+        code_challenge=code_challenge,
+        exp=int(time.time()) + 300,
+    ))
+    db.session.commit()
+
+    return redirect(f"{redirect_uri}?code={code}&state={state}")
+
+
+@oidc_blueprint.route("/token", methods=["POST"])
+@bypass_csrf_protection
+def token():
+    auth_header = request.headers.get("Authorization")
+    if not auth_header or not auth_header.startswith("Basic "):
+        return jsonify({"error": "invalid_client"}), 401
+
+    try:
+        b64 = auth_header.split(" ", 1)[1]
+        decoded = base64.b64decode(b64).decode("utf-8")
+        client_id, client_secret = decoded.split(":", 1)
+    except Exception:
+        return jsonify({"error": "invalid_client"}), 401
+
+    client = OIDCClient.query.filter_by(client_id=client_id).first()
+    if not client or client.client_secret != client_secret:
+        return jsonify({"error": "invalid_client"}), 401
+
+    issuer = get_config("base_url", "") + "/oidc/"
+
+    def create_tokens(thing):
+        user = thing.user
+        db.session.delete(thing)
+
+        now = int(time.time())
+
+        refresh_token = secrets.token_urlsafe(48)
+        db.session.add(OIDCRefreshToken(
+            refresh_token=refresh_token,
+            user_id=user.id,
+            client_id=client.client_id,
+            exp=now + 86400,
+        ))
+
+        access_token = secrets.token_urlsafe(48)
+        db.session.add(OIDCAccessToken(
+            access_token=access_token,
+            user_id=user.id,
+            client_id=client.client_id,
+            exp=now + 3600,
+        ))
+
+        db.session.commit()
+
+        id_token = sign_jwt(
+            {"sub": str(user.id), "email": user.email, "name": user.name},
+            issuer,
+            client.client_id,
+        )
+
+        return jsonify({
+            "access_token": access_token,
+            "id_token": id_token,
+            "refresh_token": refresh_token,
+            "token_type": "Bearer",
+            "expires_in": 3600,
+        })
+
+    if request.form["grant_type"] == "authorization_code":
+        code = request.form.get("code")
+        redirect_uri = request.form.get("redirect_uri")
+
+        authcode = OIDCAuthCode.query.filter_by(code=code, client_id=client.client_id, redirect_uri=redirect_uri).first()
+        if not authcode or authcode.exp < time.time():
+            return jsonify({"error": "invalid_grant"}), 400
+
+        if client.pkce:
+            verifier = request.form.get("code_verifier")
+            digest = hashlib.sha256(verifier.encode()).digest()
+            challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
+
+            if challenge != authcode.code_challenge:
+                return jsonify({"error": "invalid_grant"}), 400
+
+        return create_tokens(authcode)
+
+    if request.form["grant_type"] == "refresh_token":
+        old_refresh = request.form["refresh_token"]
+
+        token = OIDCRefreshToken.query.filter_by(refresh_token=old_refresh, client_id=client.client_id).first()
+
+        if (not token or token.exp < time.time()):
+            return jsonify({"error": "invalid_grant"}), 400
+
+        return create_tokens(token)
+
+    return jsonify({"error": "unsupported_grant_type"}), 400
+
+
+@oidc_blueprint.route("/userinfo")
+def userinfo():
+    auth = request.headers.get("Authorization", "")
+    if not auth.startswith("Bearer "):
+        return jsonify({"error": "invalid_token"}), 401
+
+    access_token = auth.split(" ", 1)[1]
+
+    token = OIDCAccessToken.query.filter_by(access_token=access_token).first()
+
+    if not token or token.exp < time.time():
+        return jsonify({"error": "invalid_token"}), 401
+
+    user = Users.query.get(token.user_id)
+    if not user:
+        return jsonify({"error": "invalid_token"}), 401
+
+    team_id = None
+    team_name = None
+
+    if hasattr(user, "team") and user.team:
+        team_id = str(user.team.id)
+        team_name = user.team.name
+
+    return jsonify({
+        "sub": str(user.id),
+        "email": user.email,
+        "name": user.name,
+        "team_id": team_id,
+        "team_name": team_name,
+    })
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..942a25e
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,4 @@
+cffi==2.0.0
+cryptography==46.0.3
+pycparser==2.23
+PyJWT==2.10.1
diff --git a/templates/index.html b/templates/index.html
new file mode 100644
index 0000000..fa6c165
--- /dev/null
+++ b/templates/index.html
@@ -0,0 +1,112 @@
+{% extends "admin/base.html" %}
+{% block content %}
+<div class="jumbotron">
+    <div class="container">
+        <h1>
+            OIDC Identity Provider Configuration
+        </h1>
+    </div>
+</div>
+
+<div class="container">
+    <div class="row pb-4">
+        <div class="col-md-12">
+            <ul class="nav nav-tabs nav-fill" role="tablist">
+                <li class="nav-item">
+                    <a class="nav-link active" data-toggle="tab" href="#setup" role="tab">
+                        Setup
+                    </a>
+                </li>
+                <li class="nav-item">
+                    <a class="nav-link" data-toggle="tab" href="#clients" role="tab">
+                        Clients
+                    </a>
+                </li>
+                <li class="nav-item">
+                    <a class="nav-link" data-toggle="tab" href="#keys" role="tab">
+                        Signing Keys
+                    </a>
+                </li>
+            </ul>
+        </div>
+    </div>
+    <div class="row">
+        <div class="col-md-12 table-responsive">
+            <div class="tab-content">
+                <div class="tab-pane fade show active" id="setup" role="tabpanel">
+                    <form method="POST" action="{{ url_for('oidc_admin.config') }}">
+                        <input id="nonce" name='nonce' type='hidden' value="{{ Session.nonce }}">
+
+                        <div class="form-group">
+                            <b><label>Base URL</label></b>
+                            <input name="base_url" class="form-control" value="{{ baseUrl }}" required>
+                            <small class="form-text text-muted">
+                                Set this to the URL that this CTFd instance is accessible at (including the protocol).
+                                This is required for OIDC discovery to work properly.
+                            </small>
+                        </div>
+
+                        <button class="btn btn-primary mt-3">Save</button>
+                    </form>
+                </div>
+                <div class="tab-pane fade" id="clients" role="tabpanel">
+                    <table class="table pb-8">
+                    <tr>
+                        <th>Client ID</th>
+                        <th>Client secret</th>
+                        <th>Redirect URIs</th>
+                        <th></th>
+                    </tr>
+                    {% for c in clients %}
+                    <tr>
+                        <td>{{ c.client_id }}</td>
+                        <td>{{ c.client_secret }}</td>
+                        <td><small>{{ c.redirect_uris }}</small></td>
+                        <td>
+                        <form method="POST" action="{{ url_for('oidc_admin.delete_client', client_id=c.client_id) }}">
+                            <input id="nonce" name='nonce' type='hidden' value="{{ Session.nonce }}">
+
+                            <button class="btn btn-danger btn-sm">Delete</button>
+                        </form>
+                        </td>
+                    </tr>
+                    {% endfor %}
+                    </table>
+
+                    <h2>Create new client</h2>
+
+                    <form method="POST" action="{{ url_for('oidc_admin.clients' )}}">
+                        <input id="nonce" name='nonce' type='hidden' value="{{ Session.nonce }}">
+
+                        <div class="form-group">
+                            <b><label>Client ID</label></b>
+                            <input name="client_id" class="form-control" required>
+                        </div>
+
+                        <div class="form-group">
+                            <b><label>Redirect URIs (one per line)</label></b>
+                            <textarea name="redirect_uris" class="form-control" required></textarea>
+                        </div>
+
+                        <button class="btn btn-primary mt-3">Create</button>
+                    </form>
+                </div>
+                <div class="tab-pane fade" id="keys" role="tabpanel">
+                    <table class="table pb-4">
+                    <tr>
+                        <th>KID</th>
+                        <th>Created</th>
+                    </tr>
+                    {% for k in keys %}
+                    <tr>
+                        <td>{{ k.kid }}</td>
+                        <td>{{ k.created }}</td>
+                    </tr>
+                    {% endfor %}
+                    </table>
+                </div>
+            </div>
+        </div>
+    </div>
+</div>
+{% endblock %}