munch-ease-backend/security.py

from __future__ import annotations

import base64
import hmac
import json
import os
import time
from dataclasses import dataclass
from hashlib import sha256
from typing import Any, Dict, Tuple


def _b64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_decode(data: str) -> bytes:
    padding = "=" * (-len(data) % 4)
    return base64.urlsafe_b64decode(data + padding)


@dataclass
class JwtConfig:
    issuer: str
    audience: str
    access_secret: bytes
    refresh_secret: bytes
    access_ttl_seconds: int
    refresh_ttl_seconds: int


def _sign(secret: bytes, msg: bytes) -> str:
    sig = hmac.new(secret, msg, sha256).digest()
    return _b64url_encode(sig)


def _encode_header() -> str:
    header = {"alg": "HS256", "typ": "JWT"}
    return _b64url_encode(json.dumps(header, separators=(",", ":")).encode("utf-8"))


def _encode_payload(claims: Dict[str, Any]) -> str:
    return _b64url_encode(json.dumps(claims, separators=(",", ":")).encode("utf-8"))


def create_jwt(
    config: JwtConfig, subject: str, kind: str = "access", extra: Dict[str, Any] | None = None
) -> str:
    now = int(time.time())
    ttl = config.access_ttl_seconds if kind == "access" else config.refresh_ttl_seconds
    secret = config.access_secret if kind == "access" else config.refresh_secret
    claims: Dict[str, Any] = {
        "iss": config.issuer,
        "aud": config.audience,
        "sub": subject,
        "iat": now,
        "exp": now + ttl,
        "typ": kind,
    }
    if extra:
        claims.update(extra)
    header = _encode_header()
    payload = _encode_payload(claims)
    signing_input = f"{header}.{payload}".encode("ascii")
    signature = _sign(secret, signing_input)
    return f"{header}.{payload}.{signature}"


def verify_jwt(
    config: JwtConfig, token: str, expected_kind: str = "access"
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    try:
        header_b64, payload_b64, sig = token.split(".")
    except ValueError:
        raise ValueError("Invalid token format")
    signing_input = f"{header_b64}.{payload_b64}".encode("ascii")
    header = json.loads(_b64url_decode(header_b64))
    if header.get("alg") != "HS256" or header.get("typ") != "JWT":
        raise ValueError("Unsupported JWT header")
    payload = json.loads(_b64url_decode(payload_b64))
    kind = payload.get("typ")
    secret = config.access_secret if kind == "access" else config.refresh_secret
    if not hmac.compare_digest(sig, _sign(secret, signing_input)):
        raise ValueError("Invalid signature")
    now = int(time.time())
    if payload.get("iss") != config.issuer or payload.get("aud") != config.audience:
        raise ValueError("Invalid claims")
    if kind != expected_kind:
        raise ValueError("Invalid token type")
    if int(payload.get("exp", 0)) < now:
        raise ValueError("Token expired")
    return header, payload


def random_secret(n: int = 32) -> bytes:
    return os.urandom(n)