casproxy.py

import hashlib
import hmac
import importlib
import json
import logging
import os
import re
import sqlite3
import time
import uuid
from dataclasses import dataclass
from typing import Any, Optional
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

import jwt
import requests
from defusedxml import ElementTree
from flask import Flask, jsonify, redirect, render_template, request
from pypinyin import lazy_pinyin


CAS_NS = {"cas": "http://www.yale.edu/tp/cas"}
DEFAULT_CLEANUP_INTERVAL = 600


@dataclass(frozen=True)
class Settings:
    cas_login_url: str
    cas_validate_url: str
    service_url: str
    redirect_uri_allowlist: tuple[str, ...]
    client_id: str
    client_secret: str
    jwt_secret_key: str
    access_token_expiry: int = 1800
    auth_code_expiry: int = 300
    session_expiry: int = 600
    database_path: str = "cas_proxy.sqlite3"
    issuer: str = "cas-proxy"
    cas_request_timeout: float = 8.0
    cookie_secure: bool = True
    cookie_samesite: str = "Lax"
    cleanup_interval: int = DEFAULT_CLEANUP_INTERVAL
    host: str = "127.0.0.1"
    port: int = 59084
    require_token_redirect_uri: bool = False


def _load_legacy_config() -> Any:
    try:
        return importlib.import_module("config")
    except ModuleNotFoundError:
        return None


def _get_config_value(config_module: Any, name: str, default: Any = None) -> Any:
    return os.getenv(name, getattr(config_module, name, default))


def _as_bool(value: Any) -> bool:
    if isinstance(value, bool):
        return value
    return str(value).strip().lower() in {"1", "true", "yes", "on"}


def _as_int(value: Any, default: int) -> int:
    if value is None or value == "":
        return default
    return int(value)


def _as_float(value: Any, default: float) -> float:
    if value is None or value == "":
        return default
    return float(value)


def _as_allowlist(value: Any) -> tuple[str, ...]:
    if not value:
        return ()
    if isinstance(value, (list, tuple, set)):
        return tuple(str(item).strip() for item in value if str(item).strip())
    return tuple(item.strip() for item in str(value).split(",") if item.strip())


def load_settings() -> Settings:
    config_module = _load_legacy_config()

    service_url = _get_config_value(config_module, "SERVICE_URL", "")
    redirect_allowlist = _get_config_value(
        config_module,
        "REDIRECT_URI_ALLOWLIST",
        _get_config_value(config_module, "REDIRECT_URL_MATCH", ""),
    )

    required = {
        "CAS_LOGIN_URL": _get_config_value(config_module, "CAS_LOGIN_URL", ""),
        "CAS_VALIDATE_URL": _get_config_value(config_module, "CAS_VALIDATE_URL", ""),
        "SERVICE_URL": service_url,
        "CLIENT_ID": _get_config_value(config_module, "CLIENT_ID", ""),
        "CLIENT_SECRET": _get_config_value(config_module, "CLIENT_SECRET", ""),
        "SECRET_KEY": _get_config_value(config_module, "SECRET_KEY", ""),
    }
    missing = [name for name, value in required.items() if not value]
    if not _as_allowlist(redirect_allowlist):
        missing.append("REDIRECT_URI_ALLOWLIST or REDIRECT_URL_MATCH")
    if missing:
        raise RuntimeError(f"Missing required config values: {', '.join(missing)}")

    cookie_secure_default = urlsplit(service_url).scheme == "https"

    return Settings(
        cas_login_url=required["CAS_LOGIN_URL"],
        cas_validate_url=required["CAS_VALIDATE_URL"],
        service_url=service_url,
        redirect_uri_allowlist=_as_allowlist(redirect_allowlist),
        client_id=required["CLIENT_ID"],
        client_secret=required["CLIENT_SECRET"],
        jwt_secret_key=required["SECRET_KEY"],
        access_token_expiry=_as_int(
            _get_config_value(config_module, "ACCESS_TOKEN_EXPIRY", 1800), 1800
        ),
        auth_code_expiry=_as_int(
            _get_config_value(config_module, "AUTH_CODE_EXPIRY", 300), 300
        ),
        session_expiry=_as_int(
            _get_config_value(config_module, "SESSION_EXPIRY", 600), 600
        ),
        database_path=_get_config_value(
            config_module, "DATABASE_PATH", "cas_proxy.sqlite3"
        ),
        issuer=_get_config_value(config_module, "ISSUER", "cas-proxy"),
        cas_request_timeout=_as_float(
            _get_config_value(config_module, "CAS_REQUEST_TIMEOUT", 8.0), 8.0
        ),
        cookie_secure=_as_bool(
            _get_config_value(config_module, "COOKIE_SECURE", cookie_secure_default)
        ),
        cookie_samesite=_get_config_value(config_module, "COOKIE_SAMESITE", "Lax"),
        cleanup_interval=_as_int(
            _get_config_value(
                config_module, "CLEANUP_INTERVAL", DEFAULT_CLEANUP_INTERVAL
            ),
            DEFAULT_CLEANUP_INTERVAL,
        ),
        host=_get_config_value(config_module, "HOST", "127.0.0.1"),
        port=_as_int(_get_config_value(config_module, "PORT", 59084), 59084),
        require_token_redirect_uri=_as_bool(
            _get_config_value(config_module, "REQUIRE_TOKEN_REDIRECT_URI", False)
        ),
    )


class StateStore:
    def __init__(self, database_path: str):
        self.database_path = database_path
        self._last_cleanup = 0
        self.init_db()

    def connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(
            self.database_path,
            timeout=10,
            isolation_level=None,
            check_same_thread=False,
        )
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA busy_timeout=5000")
        return conn

    def init_db(self) -> None:
        with self.connect() as conn:
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    redirect_uri TEXT NOT NULL,
                    state TEXT,
                    nonce TEXT,
                    expires_at INTEGER NOT NULL
                );

                CREATE TABLE IF NOT EXISTS auth_codes (
                    code TEXT PRIMARY KEY,
                    user_data TEXT NOT NULL,
                    redirect_uri TEXT NOT NULL,
                    expires_at INTEGER NOT NULL
                );

                CREATE TABLE IF NOT EXISTS tokens (
                    token_hash TEXT PRIMARY KEY,
                    user_data TEXT NOT NULL,
                    expires_at INTEGER NOT NULL
                );

                CREATE INDEX IF NOT EXISTS idx_sessions_expires_at
                    ON sessions(expires_at);
                CREATE INDEX IF NOT EXISTS idx_auth_codes_expires_at
                    ON auth_codes(expires_at);
                CREATE INDEX IF NOT EXISTS idx_tokens_expires_at
                    ON tokens(expires_at);
                """
            )

    def cleanup_expired(self, now: Optional[int] = None) -> None:
        now = now or int(time.time())
        with self.connect() as conn:
            conn.execute("DELETE FROM sessions WHERE expires_at <= ?", (now,))
            conn.execute("DELETE FROM auth_codes WHERE expires_at <= ?", (now,))
            conn.execute("DELETE FROM tokens WHERE expires_at <= ?", (now,))

    def cleanup_expired_if_due(self, interval: int) -> None:
        now = int(time.time())
        if now - self._last_cleanup < interval:
            return
        self.cleanup_expired(now)
        self._last_cleanup = now

    def create_session(
        self,
        session_id: str,
        redirect_uri: str,
        state: Optional[str],
        nonce: Optional[str],
        expires_at: int,
    ) -> None:
        with self.connect() as conn:
            conn.execute(
                """
                INSERT INTO sessions
                    (session_id, redirect_uri, state, nonce, expires_at)
                VALUES (?, ?, ?, ?, ?)
                """,
                (session_id, redirect_uri, state, nonce, expires_at),
            )

    def consume_session(self, session_id: str, now: Optional[int] = None) -> Optional[dict]:
        now = now or int(time.time())
        with self.connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            try:
                row = conn.execute(
                    """
                    SELECT session_id, redirect_uri, state, nonce, expires_at
                    FROM sessions
                    WHERE session_id = ? AND expires_at > ?
                    """,
                    (session_id, now),
                ).fetchone()
                if row:
                    conn.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
                conn.execute("COMMIT")
            except Exception:
                conn.execute("ROLLBACK")
                raise
        return dict(row) if row else None

    def create_auth_code(
        self, code: str, user_data: dict, redirect_uri: str, expires_at: int
    ) -> None:
        with self.connect() as conn:
            conn.execute(
                """
                INSERT INTO auth_codes
                    (code, user_data, redirect_uri, expires_at)
                VALUES (?, ?, ?, ?)
                """,
                (code, json.dumps(user_data), redirect_uri, expires_at),
            )

    def consume_auth_code(self, code: str, now: Optional[int] = None) -> Optional[dict]:
        now = now or int(time.time())
        with self.connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            try:
                row = conn.execute(
                    """
                    SELECT code, user_data, redirect_uri, expires_at
                    FROM auth_codes
                    WHERE code = ? AND expires_at > ?
                    """,
                    (code, now),
                ).fetchone()
                if row:
                    conn.execute("DELETE FROM auth_codes WHERE code = ?", (code,))
                conn.execute("COMMIT")
            except Exception:
                conn.execute("ROLLBACK")
                raise
        if not row:
            return None
        data = dict(row)
        data["user_data"] = json.loads(data["user_data"])
        return data

    def store_token(self, token: str, user_data: dict, expires_at: int) -> None:
        with self.connect() as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO tokens
                    (token_hash, user_data, expires_at)
                VALUES (?, ?, ?)
                """,
                (hash_token(token), json.dumps(user_data), expires_at),
            )

    def get_token_user_data(
        self, token: str, now: Optional[int] = None
    ) -> Optional[dict]:
        now = now or int(time.time())
        with self.connect() as conn:
            row = conn.execute(
                """
                SELECT user_data
                FROM tokens
                WHERE token_hash = ? AND expires_at > ?
                """,
                (hash_token(token), now),
            ).fetchone()
        return json.loads(row["user_data"]) if row else None


def hash_token(token: str) -> str:
    return hashlib.sha256(token.encode("utf-8")).hexdigest()


def split_name(name: str) -> tuple[str, str]:
    if not name:
        return "", ""

    if re.fullmatch(r"[A-Za-z ]+", name):
        parts = name.split()
        return parts[0], " ".join(parts[1:]) if len(parts) > 1 else ""

    if len(name) <= 5 and re.fullmatch(r"[\u4e00-\u9fa5]+", name):
        lastname, firstname = name[0], name[1:]
        lastname_pinyin = "".join(lazy_pinyin(lastname)).upper()
        firstname_pinyin = "".join(lazy_pinyin(firstname)).capitalize()
        return firstname_pinyin, lastname_pinyin

    return name, ""


def compare_secret(left: Optional[str], right: str) -> bool:
    if left is None:
        return False
    return hmac.compare_digest(left, right)


def is_redirect_uri_allowed(redirect_uri: str, allowlist: tuple[str, ...]) -> bool:
    if not redirect_uri or not allowlist:
        return False

    candidate = urlsplit(redirect_uri)
    if candidate.scheme not in {"http", "https"} or not candidate.netloc:
        return False

    for allowed_uri in allowlist:
        allowed = urlsplit(allowed_uri)
        if candidate.scheme != allowed.scheme:
            continue
        if (candidate.hostname or "").lower() != (allowed.hostname or "").lower():
            continue
        if candidate.port != allowed.port:
            continue

        allowed_path = allowed.path.rstrip("/")
        if allowed_path and allowed_path != "":
            if candidate.path != allowed_path and not candidate.path.startswith(
                f"{allowed_path}/"
            ):
                continue

        return True

    return False


def add_query_params(url: str, params: dict[str, Optional[str]]) -> str:
    parts = urlsplit(url)
    query = dict(parse_qsl(parts.query, keep_blank_values=True))
    query.update({key: value for key, value in params.items() if value is not None})
    return urlunsplit(
        (
            parts.scheme,
            parts.netloc,
            parts.path,
            urlencode(query),
            parts.fragment,
        )
    )


def get_basic_auth_client_credentials() -> tuple[Optional[str], Optional[str]]:
    auth = request.authorization
    if not auth:
        return None, None
    return auth.username, auth.password


def get_token_client_credentials() -> tuple[Optional[str], Optional[str]]:
    basic_client_id, basic_client_secret = get_basic_auth_client_credentials()
    if basic_client_id or basic_client_secret:
        return basic_client_id, basic_client_secret
    return request.form.get("client_id"), request.form.get("client_secret")


def cas_text(parent: Any, tag_name: str) -> Optional[str]:
    if parent is None:
        return None
    item = parent.find(f".//cas:{tag_name}", CAS_NS)
    if item is None or item.text is None:
        return None
    value = item.text.strip()
    return value or None


def parse_cas_response(xml_content: bytes) -> dict:
    tree = ElementTree.fromstring(xml_content)
    failure = tree.find(".//cas:authenticationFailure", CAS_NS)
    if failure is not None:
        reason = (failure.text or "CAS authentication failed").strip()
        raise ValueError(reason)

    user = cas_text(tree, "user")
    attributes = tree.find(".//cas:attributes", CAS_NS)
    sid = cas_text(attributes, "sid")
    name = cas_text(attributes, "name") or user
    email = cas_text(attributes, "email")

    if not user:
        raise ValueError("CAS response did not include user")
    if attributes is None:
        raise ValueError("CAS response did not include attributes")
    if not sid:
        raise ValueError("CAS response did not include sid")
    if not name:
        raise ValueError("CAS response did not include name")

    if not email and sid.startswith("1"):
        email = f"{sid}@mail.sustech.edu.cn"

    return {
        "sub": user,
        "sid": sid,
        "username": sid,
        "name": name,
        "email": email,
    }


def validate_cas_ticket(ticket: str, settings: Settings) -> dict:
    params = {"service": settings.service_url, "ticket": ticket, "format": "XML"}
    response = requests.get(
        settings.cas_validate_url,
        params=params,
        timeout=settings.cas_request_timeout,
    )
    response.raise_for_status()
    return parse_cas_response(response.content)


def create_token_claims(
    settings: Settings,
    user_data: dict,
    now: int,
    expires_at: int,
    nonce: Optional[str],
    include_profile: bool = False,
) -> dict:
    claims = {
        "iss": settings.issuer,
        "aud": settings.client_id,
        "sub": user_data["sub"],
        "iat": now,
        "nbf": now,
        "exp": expires_at,
        "jti": str(uuid.uuid4()),
    }
    if nonce:
        claims["nonce"] = nonce
    if include_profile:
        claims.update(user_data)
    return claims


def encode_jwt(claims: dict, secret: str) -> str:
    token = jwt.encode(claims, secret, algorithm="HS256")
    if isinstance(token, bytes):
        return token.decode("utf-8")
    return token


def create_app(settings: Optional[Settings] = None) -> Flask:
    settings = settings or load_settings()
    app = Flask(__name__)
    app.config["CAS_PROXY_SETTINGS"] = settings
    app.config["CAS_PROXY_STORE"] = StateStore(settings.database_path)

    @app.before_request
    def cleanup_expired_records() -> None:
        store: StateStore = app.config["CAS_PROXY_STORE"]
        store.cleanup_expired_if_due(settings.cleanup_interval)

    @app.route("/authorize")
    def authorize():
        client_id = request.args.get("client_id")
        if not compare_secret(client_id, settings.client_id):
            return jsonify(error="invalid_client"), 400

        response_type = request.args.get("response_type", "code")
        if response_type != "code":
            return jsonify(error="unsupported_response_type"), 400

        redirect_uri = request.args.get("redirect_uri")
        if not is_redirect_uri_allowed(redirect_uri or "", settings.redirect_uri_allowlist):
            return jsonify(error="invalid_redirect_uri"), 400

        state = request.args.get("state")
        nonce = request.args.get("nonce")
        session_id = f"SESSION-{uuid.uuid4()}"
        now = int(time.time())

        store: StateStore = app.config["CAS_PROXY_STORE"]
        store.create_session(
            session_id=session_id,
            redirect_uri=redirect_uri,
            state=state,
            nonce=nonce,
            expires_at=now + settings.session_expiry,
        )

        login_url = add_query_params(settings.cas_login_url, {"service": settings.service_url})
        response = redirect(login_url)
        response.set_cookie(
            "session_id",
            session_id,
            max_age=settings.session_expiry,
            httponly=True,
            secure=settings.cookie_secure,
            samesite=settings.cookie_samesite,
        )
        return response

    @app.route("/callback")
    def callback():
        session_id = request.cookies.get("session_id")
        if not session_id:
            return jsonify(error="missing_session"), 400

        ticket = request.args.get("ticket")
        if not ticket:
            return jsonify(error="missing_ticket"), 400

        store: StateStore = app.config["CAS_PROXY_STORE"]
        session_data = store.consume_session(session_id)
        if not session_data:
            return jsonify(error="invalid_or_expired_session"), 400

        try:
            cas_user = validate_cas_ticket(ticket, settings)
        except requests.RequestException:
            app.logger.exception("Failed to validate CAS ticket")
            return jsonify(error="cas_validation_failed"), 502
        except (ElementTree.ParseError, ValueError):
            app.logger.exception("Invalid CAS response")
            return jsonify(error="invalid_cas_response"), 400

        if not cas_user["email"]:
            response = render_template("error-cas-email-not-found.html")
            return response, 400

        first_name, last_name = split_name(cas_user["name"])
        user_data = {
            "sub": cas_user["sub"],
            "sid": cas_user["sid"],
            "username": cas_user["username"],
            "name": cas_user["name"],
            "email": cas_user["email"],
            "given_name": first_name,
            "family_name": last_name,
            "nonce": session_data["nonce"],
        }

        auth_code = str(uuid.uuid4())
        now = int(time.time())
        store.create_auth_code(
            code=auth_code,
            user_data=user_data,
            redirect_uri=session_data["redirect_uri"],
            expires_at=now + settings.auth_code_expiry,
        )

        response_url = add_query_params(
            session_data["redirect_uri"],
            {
                "code": auth_code,
                "state": session_data["state"],
                "nonce": session_data["nonce"],
            },
        )
        response = redirect(response_url)
        response.delete_cookie("session_id")
        return response

    @app.route("/token", methods=["POST"])
    def token():
        client_id, client_secret = get_token_client_credentials()
        if not compare_secret(client_id, settings.client_id) or not compare_secret(
            client_secret, settings.client_secret
        ):
            return jsonify(error="invalid_client"), 401

        grant_type = request.form.get("grant_type")
        if grant_type and grant_type != "authorization_code":
            return jsonify(error="unsupported_grant_type"), 400

        auth_code = request.form.get("code")
        if not auth_code:
            return jsonify(error="invalid_request", error_description="Missing code"), 400

        store: StateStore = app.config["CAS_PROXY_STORE"]
        auth_data = store.consume_auth_code(auth_code)
        if not auth_data:
            return jsonify(error="invalid_grant"), 400

        token_redirect_uri = request.form.get("redirect_uri")
        if token_redirect_uri:
            if token_redirect_uri != auth_data["redirect_uri"]:
                return jsonify(error="invalid_grant"), 400
        elif settings.require_token_redirect_uri:
            return jsonify(error="invalid_request", error_description="Missing redirect_uri"), 400

        now = int(time.time())
        expires_at = now + settings.access_token_expiry
        user_data = auth_data["user_data"]

        access_token = encode_jwt(
            create_token_claims(
                settings,
                user_data,
                now,
                expires_at,
                nonce=None,
                include_profile=False,
            ),
            settings.jwt_secret_key,
        )
        id_token = encode_jwt(
            create_token_claims(
                settings,
                user_data,
                now,
                expires_at,
                nonce=user_data.get("nonce"),
                include_profile=True,
            ),
            settings.jwt_secret_key,
        )
        store.store_token(access_token, user_data, expires_at)

        return jsonify(
            access_token=access_token,
            id_token=id_token,
            token_type="Bearer",
            scope="openid profile email",
            expires_in=settings.access_token_expiry,
        )

    @app.route("/userinfo")
    def userinfo():
        auth_header = request.headers.get("Authorization", "")
        parts = auth_header.split()
        if len(parts) != 2 or parts[0].lower() != "bearer":
            return jsonify(error="invalid_request"), 401

        access_token = parts[1]
        try:
            jwt.decode(
                access_token,
                settings.jwt_secret_key,
                algorithms=["HS256"],
                audience=settings.client_id,
                issuer=settings.issuer,
            )
        except jwt.InvalidTokenError:
            return jsonify(error="invalid_token"), 401

        store: StateStore = app.config["CAS_PROXY_STORE"]
        user_data = store.get_token_user_data(access_token)
        if not user_data:
            return jsonify(error="invalid_token"), 401

        return jsonify(user_data)

    @app.route("/healthz")
    def healthz():
        return jsonify(status="ok")

    return app


try:
    app = create_app()
except RuntimeError as exc:
    logging.getLogger(__name__).warning("CAS proxy is not fully configured: %s", exc)
    app = Flask(__name__)

    @app.route("/healthz")
    def healthz_unconfigured():
        return jsonify(status="unconfigured"), 500

    @app.route("/", defaults={"path": ""})
    @app.route("/<path:path>")
    def unconfigured(path: str):
        return jsonify(error="server_not_configured"), 500


if __name__ == "__main__":
    settings = load_settings()
    app = create_app(settings)
    app.run(host=settings.host, port=settings.port, debug=False, threaded=True)