app/backend/fluksio/core/security.py

import base64
import uuid
from datetime import UTC, datetime, timedelta
from functools import lru_cache
from typing import Any

import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pwdlib import PasswordHash
from pwdlib.hashers.argon2 import Argon2Hasher
from pwdlib.hashers.bcrypt import BcryptHasher

from fluksio.core.config import settings

password_hash = PasswordHash(
    (
        Argon2Hasher(),
        BcryptHasher(),
    )
)


ALGORITHM = "HS256"
#: MCP tokens are signed with a keypair of their own, so the public half can be
#: published and the whole set revoked by rotating it — without logging anyone
#: out of the browser, and without the resource server needing the secret that
#: signs browser sessions.
OAUTH_ALGORITHM = "RS256"
#: The one scope an MCP token carries.
MCP_SCOPE = "mcp"
#: What a remote worker's credential says it is for. Its own audience, so an
#: agent's token cannot attach a worker and a worker's cannot call the API.
WORKER_SCOPE = "worker"
WORKER_AUDIENCE = "fluksio-worker"
#: What a paired wall panel's credential says it is for. Its own audience, so
#: the session decode refuses it outright and the scope check in ``deps`` is
#: the only door it fits.
PANEL_AUDIENCE = "fluksio-panel"


def create_access_token(subject: str | Any, expires_delta: timedelta) -> str:
    expire = datetime.now(UTC) + expires_delta
    to_encode = {"exp": expire, "sub": str(subject)}
    encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt


# -----------------------------------------------------------------------------
# OAuth signing key
# -----------------------------------------------------------------------------


def _load_or_create_key() -> rsa.RSAPrivateKey:
    """The RSA key MCP tokens are signed with, generated on first use."""
    path = settings.OAUTH_PRIVATE_KEY_FILE
    if path.exists():
        loaded = serialization.load_pem_private_key(path.read_bytes(), password=None)
        if not isinstance(loaded, rsa.RSAPrivateKey):
            raise TypeError(f"{path} is not an RSA private key")
        return loaded

    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(
        key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
    )
    path.chmod(0o600)
    return key


@lru_cache(maxsize=1)
def oauth_key() -> rsa.RSAPrivateKey:
    return _load_or_create_key()


def public_jwks() -> dict[str, Any]:
    """The public half, for anything that wants to check a token itself."""
    numbers = oauth_key().public_key().public_numbers()

    def b64(value: int) -> str:
        raw = value.to_bytes((value.bit_length() + 7) // 8, "big")
        return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")

    return {
        "keys": [
            {
                "kty": "RSA",
                "use": "sig",
                "alg": OAUTH_ALGORITHM,
                "kid": "fluksio-oauth",
                "n": b64(numbers.n),
                "e": b64(numbers.e),
            }
        ]
    }


def create_oauth_access_token(
    user_id: uuid.UUID, client_id: uuid.UUID, expires_delta: timedelta
) -> str:
    """An access token for the MCP channel, told apart from a browser session.

    The ``mcp`` claim is what the MCP endpoint checks: a perfectly valid token
    from someone's browser is refused there, so agent traffic never arrives
    looking like a person's.
    """
    now = datetime.now(UTC)
    payload = {
        "sub": str(user_id),
        "iss": settings.oauth_issuer,
        "aud": settings.mcp_resource,
        "iat": now,
        "exp": now + expires_delta,
        "mcp": True,
        "client_id": str(client_id),
        "scope": MCP_SCOPE,
    }
    return jwt.encode(
        payload,
        oauth_key().private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        ),
        algorithm=OAUTH_ALGORITHM,
        headers={"kid": "fluksio-oauth"},
    )


def create_worker_token(
    name: str,
    expires_delta: timedelta,
    labels: list[str] | None = None,
    jti: str | None = None,
) -> str:
    """A credential a remote worker presents when it dials in.

    Signed with the same keypair the agent tokens use, and told apart from them
    by its audience: a worker's token grants no API access, and an agent's
    opens no worker connection. It names the worker and the labels it may
    carry, so a worker cannot claim work meant for another machine, and its
    `jti` is what revoking one token rather than the whole key goes by.
    """
    now = datetime.now(UTC)
    payload = {
        "sub": name,
        "iss": settings.oauth_issuer,
        "aud": WORKER_AUDIENCE,
        "iat": now,
        "exp": now + expires_delta,
        "scope": WORKER_SCOPE,
        "labels": sorted(set(labels or [])),
        "jti": jti or uuid.uuid4().hex,
    }
    return jwt.encode(
        payload,
        oauth_key().private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        ),
        algorithm=OAUTH_ALGORITHM,
        headers={"kid": "fluksio-oauth"},
    )


def decode_worker_token(token: str) -> dict[str, Any]:
    """Validate a worker credential. Raises ``InvalidTokenError`` if it does not."""
    payload: dict[str, Any] = jwt.decode(
        token,
        oauth_key().public_key(),
        algorithms=[OAUTH_ALGORITHM],
        audience=WORKER_AUDIENCE,
        issuer=settings.oauth_issuer,
    )
    return payload


def create_panel_token(
    panel: str, user_id: uuid.UUID | str, expires_delta: timedelta, nonce: int = 0
) -> str:
    """The credential a paired wall panel holds.

    Signed with the app's own secret like a browser session, and told apart
    from one by its audience. ``sub`` records the account that approved the
    pairing. The panel acts as itself rather than as that account (see
    ``deps.device_principal``), so deleting the account leaves it paired. The
    ``panel`` claim is what the request filter in ``fluksio.api.deps`` holds
    it to: that panel's dashboards and the message endpoints its widgets need.

    ``pnc`` is the panel's nonce at the moment of pairing, and the filter
    refuses a credential naming any other. Bumping the panel's nonce is
    therefore how one screen is re-paired without deleting the panel out from
    under its dashboards.

    Long-lived on purpose: a wall tablet is set up once and left running, and
    it has no keyboard to log in again with.
    """
    now = datetime.now(UTC)
    payload = {
        "sub": str(user_id),
        "aud": PANEL_AUDIENCE,
        "panel": panel,
        "pnc": nonce,
        "iat": now,
        "exp": now + expires_delta,
    }
    return jwt.encode(payload, settings.SECRET_KEY, algorithm=ALGORITHM)


def decode_panel_token(token: str) -> dict[str, Any]:
    """Validate a panel credential. Raises ``InvalidTokenError`` if it does not."""
    payload: dict[str, Any] = jwt.decode(
        token,
        settings.SECRET_KEY,
        algorithms=[ALGORITHM],
        audience=PANEL_AUDIENCE,
    )
    return payload


def decode_oauth_token(token: str) -> dict[str, Any]:
    """Validate an MCP token. Raises ``InvalidTokenError`` if it does not hold."""
    payload: dict[str, Any] = jwt.decode(
        token,
        oauth_key().public_key(),
        algorithms=[OAUTH_ALGORITHM],
        audience=settings.mcp_resource,
        issuer=settings.oauth_issuer,
    )
    return payload


def verify_password(
    plain_password: str, hashed_password: str
) -> tuple[bool, str | None]:
    return password_hash.verify_and_update(plain_password, hashed_password)


def get_password_hash(password: str) -> str:
    return password_hash.hash(password)