# jwt_validation.py
# coding: utf8

"""    
Décodage et validation d'un Jeton d'Identité JSON Web Token (JWT).

Auteur : 
    B.Degoy bertrand@degoy.com    
copyright (c) 2026 B.Degoy
licence : MIT
"""

import json
import base64
import hmac
import hashlib
from typing import Any, Dict


class JWTDecodeError(Exception):
    """Base class for all JWT decoding errors."""
    pass


class JWTFormatError(JWTDecodeError):
    """Raised when the JWT structure is invalid."""
    pass


class JWTAlgorithmError(JWTDecodeError):
    """Raised when the algorithm is missing or not allowed."""
    pass


class JWTSignatureError(JWTDecodeError):
    """Raised when the signature is invalid."""
    pass


class JWTUtils:
    """
    A static utility class providing low-level JWT decoding and signature
    verification, equivalent to the provided PHP implementation but using
    Python exceptions instead of returning False.

    This class does NOT validate claims (exp, nbf, iss, etc.).
    """

    # ----------------------------------------------------------------------
    # Base64URL decoding
    # ----------------------------------------------------------------------
    @staticmethod
    def urlsafe_b64decode(data: str) -> bytes:
        """
        Decode a Base64URL string into raw bytes.
        Adds required padding if missing.
        """
        padding = "=" * (-len(data) % 4)
        try:
            return base64.urlsafe_b64decode(data + padding)
        except Exception as e:
            raise JWTFormatError(f"Invalid Base64URL segment: {data}") from e

    # ----------------------------------------------------------------------
    # HMAC signing
    # ----------------------------------------------------------------------
    @staticmethod
    def sign(message: str, key: bytes, algo: str) -> bytes:
        """
        Compute an HMAC signature for the given message using HS256/384/512.
        """
        if algo == "HS256":
            return hmac.new(key, message.encode(), hashlib.sha256).digest()
        if algo == "HS384":
            return hmac.new(key, message.encode(), hashlib.sha384).digest()
        if algo == "HS512":
            return hmac.new(key, message.encode(), hashlib.sha512).digest()

        raise JWTAlgorithmError(f"Unsupported HMAC algorithm: {algo}")

    # ----------------------------------------------------------------------
    # Constant-time comparison
    # ----------------------------------------------------------------------
    @staticmethod
    def hash_equals(a: bytes, b: bytes) -> bool:
        """Constant-time comparison to avoid timing attacks."""
        return hmac.compare_digest(a, b)

    # ----------------------------------------------------------------------
    # RSA signature verification
    # ----------------------------------------------------------------------
    @staticmethod
    def verify_rsa(signature: bytes, message: str, key_pem: str, algo: str) -> bool:
        """
        Verify an RSA signature using a PEM-encoded public key.
        Supported algorithms: RS256, RS384, RS512.
        """
        from cryptography.hazmat.primitives import serialization, hashes
        from cryptography.hazmat.primitives.asymmetric import padding

        public_key = serialization.load_pem_public_key(key_pem.encode())

        hash_algo = {
            "RS256": hashes.SHA256(),
            "RS384": hashes.SHA384(),
            "RS512": hashes.SHA512(),
        }.get(algo)

        if hash_algo is None:
            raise JWTAlgorithmError(f"Unsupported RSA algorithm: {algo}")

        try:
            public_key.verify(
                signature,
                message.encode(),
                padding.PKCS1v15(),
                hash_algo
            )
            return True
        except Exception:
            return False

    # ----------------------------------------------------------------------
    # Signature verification dispatcher
    # ----------------------------------------------------------------------
    @staticmethod
    def verify_signature(signature: bytes, input_data: str, key: Any, algo: str) -> None:
        """
        Verify the signature for the given algorithm.
        Raises JWTSignatureError on failure.
        """
        if algo in ("HS256", "HS384", "HS512"):
            key_bytes = key if isinstance(key, bytes) else key.encode()
            expected = JWTUtils.sign(input_data, key_bytes, algo)

            if not JWTUtils.hash_equals(expected, signature):
                raise JWTSignatureError("Invalid HMAC signature")
            return

        if algo in ("RS256", "RS384", "RS512"):
            if not JWTUtils.verify_rsa(signature, input_data, key, algo):
                raise JWTSignatureError("Invalid RSA signature")
            return

        raise JWTAlgorithmError(f"Unsupported or invalid signing algorithm: {algo}")

    # ----------------------------------------------------------------------
    # Main decode function
    # ----------------------------------------------------------------------
    @staticmethod
    def decode(jwt: str, key: Any = None, allowed_algorithms: Any = True) -> Dict:
        """
        Decode a JWT token, validate its structure, optionally verify its
        signature, and return the payload as a dictionary.

        Raises:
            JWTFormatError: Invalid structure or Base64URL segments.
            JWTAlgorithmError: Missing or disallowed algorithm.
            JWTSignatureError: Signature verification failed.
        """
        # Basic structure check
        if "." not in jwt:
            raise JWTFormatError("JWT must contain at least one dot")

        parts = jwt.split(".")
        if len(parts) != 3:
            raise JWTFormatError("JWT must contain exactly 3 segments")

        head_b64, payload_b64, sig_b64 = parts

        # Decode JSON header and payload
        try:
            header = json.loads(JWTUtils.urlsafe_b64decode(head_b64))
        except Exception as e:
            raise JWTFormatError("Invalid JWT header") from e

        try:
            payload = json.loads(JWTUtils.urlsafe_b64decode(payload_b64))
        except Exception as e:
            raise JWTFormatError("Invalid JWT payload") from e

        signature = JWTUtils.urlsafe_b64decode(sig_b64)

        # Algorithm checks
        if allowed_algorithms:
            if "alg" not in header:
                raise JWTAlgorithmError("Missing 'alg' in JWT header")

            algo = header["alg"]

            if isinstance(allowed_algorithms, list) and algo not in allowed_algorithms:
                raise JWTAlgorithmError(f"Algorithm '{algo}' not allowed")

            signing_input = f"{head_b64}.{payload_b64}"

            JWTUtils.verify_signature(signature, signing_input, key, algo)

        return payload

from datetime import datetime, timezone


class JWTClaimsError(Exception):
    """Raised when JWT claims validation fails."""
    pass


class JWTUtils:
    # ... (tout le reste de ta classe ici)

    @staticmethod
    def validate_claims(
        payload: dict,
        issuer: str = None,
        audience: str = None,
        leeway: int = 0
    ) -> None:
        """
        Validate standard JWT claims (exp, nbf, iss, aud).

        Parameters:
            payload (dict): The decoded JWT payload.
            issuer (str|None): Expected 'iss' value.
            audience (str|None): Expected 'aud' value.
            leeway (int): Allowed clock skew in seconds.

        Raises:
            JWTClaimsError: If any claim is invalid.
        """

        now = datetime.now(timezone.utc).timestamp()

        # --- exp: expiration time ---
        if "exp" in payload:
            if now > payload["exp"] + leeway:
                raise JWTClaimsError("Token has expired")

        # --- nbf: not before ---
        if "nbf" in payload:
            if now < payload["nbf"] - leeway:
                raise JWTClaimsError("Token is not yet valid (nbf)")

        # --- iat: issued at ---
        if "iat" in payload:
            if payload["iat"] > now + leeway:
                raise JWTClaimsError("Token issued in the future (iat)")

        # --- iss: issuer ---
        if issuer is not None:
            if payload.get("iss") != issuer:
                raise JWTClaimsError(f"Invalid issuer: {payload.get('iss')}")

        # --- aud: audience ---
        if audience is not None:
            aud = payload.get("aud")
            if isinstance(aud, list):
                if audience not in aud:
                    raise JWTClaimsError(f"Audience '{audience}' not allowed")
            else:
                if aud != audience:
                    raise JWTClaimsError(f"Invalid audience: {aud}")
</python>

{{intégration dans <span class="base64-bc809d0d" title="PGNvZGUgY2xhc3M9J3NwaXBfY29kZSBzcGlwX2NvZGVfaW5saW5lJyBkaXI9J2x0cic+ZGVjb2RlKCk8L2NvZGU+"></span>}}

On ajoute un paramètre optionnel :

<code class="python">
@staticmethod
def decode(jwt, key=None, allowed_algorithms=True, validate=False, issuer=None, audience=None):
    payload = ...  # comme avant

    if validate:
        JWTUtils.validate_claims(payload, issuer=issuer, audience=audience)

    return payload
</python>

Et on l’appelle comme ceci :

<code class="python">
payload = JWTUtils.decode(
    token,
    key=public_key,
    allowed_algorithms=["RS256"],
    validate=True,
    issuer="https://auth.example.com",
    audience="my-api"
)