app/services/jwt_service.py

from __future__ import annotations

import secrets
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any, Optional

import jwt
from jwt import PyJWTError

from app.config import get_settings

_SUPPORTED_ALGORITHMS = frozenset({"HS256", "HS384", "HS512"})


class JWTService:
    def __init__(self) -> None:
        self._settings = get_settings()

    @staticmethod
    def _generate_secret() -> str:
        return secrets.token_hex(32)

    def generate(
        self,
        *,
        subject: str,
        role: Optional[str] = None,
        permissions: Optional[list[str]] = None,
        issuer: Optional[str] = None,
        audience: Optional[str] = None,
        expiry_minutes: Optional[int] = None,
        not_before_minutes: int = 0,
        extra_claims: Optional[dict[str, Any]] = None,
        secret: Optional[str] = None,
        algorithm: Optional[str] = None,
    ) -> dict[str, Any]:
        now = datetime.now(timezone.utc)
        expire_min = expiry_minutes if expiry_minutes is not None else self._settings.jwt_default_expiry_minutes

        payload: dict[str, Any] = {
            "sub": subject,
            "iat": now,
            "exp": now + timedelta(minutes=expire_min),
            "jti": uuid.uuid4().hex,
        }

        if issuer:
            payload["iss"] = issuer
        else:
            payload["iss"] = self._settings.jwt_issuer

        if audience:
            payload["aud"] = audience

        if role:
            payload["role"] = role

        if permissions:
            payload["permissions"] = permissions

        if not_before_minutes > 0:
            payload["nbf"] = now + timedelta(minutes=not_before_minutes)

        if extra_claims:
            payload.update(extra_claims)

        used_secret = secret if secret else self._generate_secret()
        used_algorithm = algorithm if algorithm else self._settings.jwt_algorithm

        if used_algorithm not in _SUPPORTED_ALGORITHMS:
            raise ValueError(f"Unsupported algorithm '{used_algorithm}'. Supported: {sorted(_SUPPORTED_ALGORITHMS)}")

        token = jwt.encode(payload, used_secret, algorithm=used_algorithm)

        valid_minutes = expire_min
        if valid_minutes < 60:
            valid_for = f"{valid_minutes} minute{'s' if valid_minutes != 1 else ''}"
        elif valid_minutes < 1440:
            hours = valid_minutes // 60
            mins = valid_minutes % 60
            valid_for = f"{hours} hour{'s' if hours != 1 else ''}"
            if mins:
                valid_for += f" {mins} minute{'s' if mins != 1 else ''}"
        else:
            days = valid_minutes // 1440
            hrs = (valid_minutes % 1440) // 60
            valid_for = f"{days} day{'s' if days != 1 else ''}"
            if hrs:
                valid_for += f" {hrs} hour{'s' if hrs != 1 else ''}"

        return {
            "token": token,
            "claims": {
                "sub": subject,
                "iss": payload["iss"],
                "aud": audience,
                "role": role,
                "permissions": permissions,
                "jti": payload["jti"],
                "iat": payload["iat"].isoformat(),
                "exp": payload["exp"].isoformat(),
                "nbf": payload.get("nbf").isoformat() if payload.get("nbf") else None,
            },
            "expires_at": payload["exp"].isoformat(),
            "valid_for": valid_for,
            "secret": used_secret,
            "algorithm": used_algorithm,
        }

    def validate(
        self,
        token: str,
        audience: Optional[str] = None,
        secret: Optional[str] = None,
        algorithm: Optional[str] = None,
    ) -> dict[str, Any]:
        used_secret = secret if secret else self._settings.jwt_secret_key
        used_algorithms = [algorithm] if algorithm else [self._settings.jwt_algorithm]

        try:
            payload = jwt.decode(
                token,
                used_secret,
                algorithms=used_algorithms,
                audience=audience,
                issuer=self._settings.jwt_issuer,
                options={
                    "require": ["sub", "exp", "iat", "jti"],
                    "verify_signature": True,
                    "verify_exp": True,
                    "verify_iat": True,
                    "verify_aud": audience is not None,
                    "verify_iss": True,
                },
            )
            return {
                "valid": True,
                "claims": {
                    "sub": payload.get("sub"),
                    "iss": payload.get("iss"),
                    "aud": payload.get("aud"),
                    "role": payload.get("role"),
                    "permissions": payload.get("permissions"),
                    "jti": payload.get("jti"),
                    "iat": datetime.fromtimestamp(payload["iat"], tz=timezone.utc).isoformat() if payload.get("iat") else None,
                    "exp": datetime.fromtimestamp(payload["exp"], tz=timezone.utc).isoformat() if payload.get("exp") else None,
                    "nbf": datetime.fromtimestamp(payload["nbf"], tz=timezone.utc).isoformat() if payload.get("nbf") else None,
                },
                "error": None,
            }
        except jwt.ExpiredSignatureError:
            return {"valid": False, "claims": None, "error": "Token has expired."}
        except jwt.InvalidAudienceError:
            return {"valid": False, "claims": None, "error": "Token audience does not match."}
        except jwt.InvalidIssuerError:
            return {"valid": False, "claims": None, "error": "Token issuer does not match."}
        except jwt.InvalidTokenError as exc:
            return {"valid": False, "claims": None, "error": f"Invalid token: {exc}"}
        except PyJWTError as exc:
            return {"valid": False, "claims": None, "error": f"JWT validation error: {exc}"}