# SPDX-License-Identifier: Apache-2.0
"""
License management for model encryption authorization.

Supports offline license verification via RSA signature.
"""

import base64
import hashlib
import json
import os
import uuid
from dataclasses import dataclass, field
from datetime import date
from typing import List, Optional

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from cryptography.hazmat.backends import default_backend
from cryptography.exceptions import InvalidSignature

from vllm.model_security.fingerprint import (
    FingerprintPolicy, get_machine_fingerprint, verify_fingerprint
)


@dataclass
class LicenseInfo:
    """License information structure."""
    license_id: str
    customer: str
    model_key_id: str          # hex key_id matching .model_key file
    dek_encrypted: str         # base64-encoded DEK encrypted with vendor private key
    dek_hash: str              # SHA-256 hash of plaintext DEK (for integrity)
    issued_at: str             # ISO date string
    expire_at: Optional[str]   # ISO date string, None = never expires
    fingerprint_policy: str    # FingerprintPolicy value
    allowed_fingerprints: List[str] = field(default_factory=list)
    max_nodes: int = 1
    signature: str = ""        # base64-encoded RSA signature (excluded from signing payload)

    def to_dict(self, include_signature: bool = True) -> dict:
        d = {
            "license_id": self.license_id,
            "customer": self.customer,
            "model_key_id": self.model_key_id,
            "dek_encrypted": self.dek_encrypted,
            "dek_hash": self.dek_hash,
            "issued_at": self.issued_at,
            "expire_at": self.expire_at,
            "fingerprint_policy": self.fingerprint_policy,
            "allowed_fingerprints": self.allowed_fingerprints,
            "max_nodes": self.max_nodes,
        }
        if include_signature:
            d["signature"] = self.signature
        return d

    @classmethod
    def from_dict(cls, d: dict) -> "LicenseInfo":
        return cls(
            license_id=d["license_id"],
            customer=d["customer"],
            model_key_id=d["model_key_id"],
            dek_encrypted=d["dek_encrypted"],
            dek_hash=d["dek_hash"],
            issued_at=d["issued_at"],
            expire_at=d.get("expire_at"),
            fingerprint_policy=d.get("fingerprint_policy", FingerprintPolicy.NONE),
            allowed_fingerprints=d.get("allowed_fingerprints", []),
            max_nodes=d.get("max_nodes", 1),
            signature=d.get("signature", ""),
        )

    def signing_payload(self) -> bytes:
        """Get canonical bytes to sign (excludes signature field)."""
        d = self.to_dict(include_signature=False)
        return json.dumps(d, sort_keys=True, separators=(',', ':')).encode('utf-8')


class LicenseManager:
    """Creates and signs license files (vendor-side tool)."""

    def __init__(self, vendor_private_key_path: str):
        with open(vendor_private_key_path, 'rb') as f:
            self.private_key = serialization.load_pem_private_key(
                f.read(), password=None, backend=default_backend()
            )

    def create_license(
        self,
        customer: str,
        model_key_id: str,
        dek: bytes,
        vendor_secret: bytes,
        fingerprint_policy: FingerprintPolicy = FingerprintPolicy.NONE,
        allowed_fingerprints: Optional[List[str]] = None,
        expire_at: Optional[str] = None,
        max_nodes: int = 1,
    ) -> LicenseInfo:
        """
        Create and sign a license.

        The DEK is encrypted symmetrically using a key derived from
        ``vendor_secret`` + ``license_id`` via HKDF, so that the same
        ``vendor_secret`` can reconstruct the DEK at load time without
        distributing the vendor private key to customers.

        Args:
            customer: Customer/organization name
            model_key_id: Hex key_id from encrypted model
            dek: Plaintext DEK bytes
            vendor_secret: Secret bytes held by the vendor; required to
                decrypt the DEK at runtime via OfflineLicenseManager.
            fingerprint_policy: Machine binding policy
            allowed_fingerprints: List of allowed machine fingerprints
            expire_at: Expiry date as ISO string (e.g. "2027-01-01"), None = no expiry
            max_nodes: Maximum number of nodes allowed (for cluster policy)

        Returns:
            Signed LicenseInfo
        """
        # Allocate license_id first so it can be used as HKDF salt when
        # deriving the DEK encryption key — matching the decryption path in
        # OfflineLicenseManager.decrypt_dek_from_license / derive_dek_key.
        license_id = str(uuid.uuid4())
        dek_hash = hashlib.sha256(dek).hexdigest()

        dek_key = OfflineLicenseManager.derive_dek_key(vendor_secret, license_id)
        dek_cipher = OfflineLicenseManager.encrypt_dek_for_license(dek, dek_key)

        info = LicenseInfo(
            license_id=license_id,
            customer=customer,
            model_key_id=model_key_id,
            dek_encrypted=dek_cipher,
            dek_hash=dek_hash,
            issued_at=date.today().isoformat(),
            expire_at=expire_at,
            fingerprint_policy=fingerprint_policy,
            allowed_fingerprints=allowed_fingerprints or [],
            max_nodes=max_nodes,
        )

        # Sign the payload
        signature = self.private_key.sign(
            info.signing_payload(),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH,
            ),
            hashes.SHA256(),
        )
        info.signature = base64.b64encode(signature).decode()
        return info

    def save_license(self, info: LicenseInfo, path: str) -> None:
        """Save license to JSON file."""
        with open(path, 'w') as f:
            json.dump(info.to_dict(), f, indent=2)

    @staticmethod
    def generate_keypair(
        private_key_path: str,
        public_key_path: str,
        key_size: int = 3072,
        passphrase: Optional[bytes] = None,
    ) -> None:
        """Generate RSA keypair for license signing.

        Args:
            private_key_path: Path to write PEM private key
            public_key_path: Path to write PEM public key
            key_size: RSA key size in bits (default 3072 for security-sensitive use)
            passphrase: Optional passphrase to protect the private key file.
                        Strongly recommended in production.
        """
        key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=key_size,
            backend=default_backend(),
        )
        # Choose encryption algorithm for private key
        if passphrase:
            encryption_algorithm = serialization.BestAvailableEncryption(passphrase)
        else:
            import warnings
            warnings.warn(
                "Generating private key without passphrase protection. "
                "Provide a passphrase in production to protect the signing key.",
                stacklevel=2,
            )
            encryption_algorithm = serialization.NoEncryption()
        # Save private key
        with open(private_key_path, 'wb') as f:
            f.write(key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=encryption_algorithm,
            ))
        # Save public key
        with open(public_key_path, 'wb') as f:
            f.write(key.public_key().public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo,
            ))


class LicenseValidator:
    """Validates licenses at runtime (user-side)."""

    def __init__(self, vendor_public_key_path: str):
        with open(vendor_public_key_path, 'rb') as f:
            self.public_key = serialization.load_pem_public_key(
                f.read(), backend=default_backend()
            )

    def load_license(self, license_path: str) -> LicenseInfo:
        """Load license from file."""
        with open(license_path, 'r') as f:
            return LicenseInfo.from_dict(json.load(f))

    def validate(
        self,
        info: LicenseInfo,
        include_gpu_in_fingerprint: bool = False,
    ) -> None:
        """
        Validate license. Raises on any failure.

        Args:
            info: License to validate
            include_gpu_in_fingerprint: Whether GPU was included in fingerprint

        Raises:
            ValueError: If license is invalid, expired, or machine doesn't match
        """
        # 1. Verify RSA signature
        self._verify_signature(info)

        # 2. Check expiry
        self._check_expiry(info)

        # 3. Check machine fingerprint
        self._check_fingerprint(info, include_gpu_in_fingerprint)

    def validate_and_extract_dek(
        self,
        license_path: str,
        model_dir: str,
        vendor_secret: Optional[bytes] = None,
        include_gpu_in_fingerprint: bool = False,
    ) -> bytes:
        """
        Validate license and extract DEK in one call.

        Args:
            license_path: Path to license file
            model_dir: Path to model directory containing .model_key
            vendor_secret: Vendor secret for DEK decryption (symmetric mode)
            include_gpu_in_fingerprint: Whether GPU was included in fingerprint

        Returns:
            Plaintext DEK bytes

        Raises:
            ValueError: If license is invalid or DEK cannot be extracted
        """
        info = self.load_license(license_path)
        self.validate(info, include_gpu_in_fingerprint)
        return self.extract_dek_from_model_dir(info, model_dir, vendor_secret)

    def extract_dek_from_model_dir(
        self,
        info: LicenseInfo,
        model_dir: str,
        vendor_secret: Optional[bytes] = None,
    ) -> bytes:
        """
        Extract DEK from model directory using license info.

        Args:
            info: Validated license info
            model_dir: Path to model directory containing .model_key
            vendor_secret: Vendor secret for DEK decryption (if using symmetric mode)

        Returns:
            Plaintext DEK bytes
        """
        model_key_path = os.path.join(model_dir, ".model_key")
        if not os.path.isfile(model_key_path):
            raise ValueError(
                f"Model key file not found: {model_key_path}. "
                "Encrypted model must include .model_key file."
            )

        with open(model_key_path, 'r') as f:
            key_data = json.load(f)

        # Verify key_id matches
        if key_data.get("key_id") != info.model_key_id:
            raise ValueError(
                "License key_id mismatch. This license is not for this model."
            )

        # Extract and decrypt DEK
        if vendor_secret:
            # Symmetric encryption mode (offline): DEK is encrypted with vendor_secret-derived key
            dek_key = OfflineLicenseManager.derive_dek_key(
                vendor_secret, info.license_id
            )
            dek = OfflineLicenseManager.decrypt_dek_from_license(
                info.dek_encrypted, dek_key
            )
        else:
            # Plaintext DEK storage is no longer supported as it provides no security.
            # Require vendor_secret for symmetric mode. For production deployments,
            # use a proper key management service or asymmetric key exchange with
            # customer-specific public keys.
            raise ValueError(
                "Cannot extract DEK: vendor_secret is required for symmetric decryption mode. "
                "Plaintext DEK storage has been removed for security. "
                "Provide vendor_secret or use a proper key management service."
            )

        # Verify DEK hash
        dek_hash = hashlib.sha256(dek).hexdigest()
        if dek_hash != info.dek_hash:
            raise ValueError(
                "DEK integrity check failed. License or model may be corrupted."
            )

        return dek

    def _verify_signature(self, info: LicenseInfo) -> None:
        """Verify RSA signature on license payload."""
        try:
            self.public_key.verify(
                base64.b64decode(info.signature),
                info.signing_payload(),
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH,
                ),
                hashes.SHA256(),
            )
        except InvalidSignature:
            raise ValueError("License signature verification failed. License may be tampered.")

    def _check_expiry(self, info: LicenseInfo) -> None:
        """Check if license is expired."""
        if info.expire_at is None:
            return
        expire_date = date.fromisoformat(info.expire_at)
        if date.today() > expire_date:
            raise ValueError(
                f"License expired on {info.expire_at}. "
                "Please contact your vendor for a license renewal."
            )

    def check_expiry_with_grace(self, info: LicenseInfo, grace_days: int = 3) -> str:
        """
        Check license expiry, returning status string instead of raising.

        Returns:
            "valid"   — not expired
            "warning" — expired but within grace period
            "expired" — expired beyond grace period, service must stop
        """
        if info.expire_at is None:
            return "valid"
        expire_date = date.fromisoformat(info.expire_at)
        today = date.today()
        if today <= expire_date:
            return "valid"
        days_over = (today - expire_date).days
        if days_over <= grace_days:
            return "warning"
        return "expired"

    def _check_fingerprint(
        self,
        info: LicenseInfo,
        include_gpu: bool = False,
    ) -> None:
        """Check machine fingerprint against license."""
        policy = FingerprintPolicy(info.fingerprint_policy)

        if policy == FingerprintPolicy.NONE:
            return

        current_fp = get_machine_fingerprint(include_gpu=include_gpu)

        if policy == FingerprintPolicy.SINGLE:
            if not info.allowed_fingerprints:
                raise ValueError("License has SINGLE policy but no allowed fingerprints.")
            if current_fp not in info.allowed_fingerprints:
                raise ValueError(
                    "Machine fingerprint mismatch. "
                    "This license is not authorized for this machine."
                )

        elif policy == FingerprintPolicy.CLUSTER:
            if not info.allowed_fingerprints:
                raise ValueError("License has CLUSTER policy but no allowed fingerprints.")
            if current_fp not in info.allowed_fingerprints:
                raise ValueError(
                    "Machine fingerprint mismatch. "
                    "This machine is not in the authorized cluster."
                )


class OfflineLicenseManager:
    """
    Simplified offline license manager that stores DEK in the model directory.

    In this approach:
    - Vendor creates encrypted model + .model_key file
    - .model_key file stores DEK encrypted with vendor's private key
    - License file stores metadata and signature
    - At runtime: validate license signature → read .model_key → decrypt DEK

    The DEK in .model_key is encrypted with vendor's public key (OAEP),
    and can only be decrypted with the vendor's private key.
    BUT: the license is issued per-customer, and the DEK is stored
    in the license itself, encrypted with customer-specific info.

    For simplicity in offline mode, we use a symmetric approach:
    - DEK is derived from license_id + vendor_secret using HKDF
    - License contains the encrypted DEK (AES-256-GCM with key=vendor_secret)
    """

    @staticmethod
    def derive_dek_key(vendor_secret: bytes, license_id: str) -> bytes:
        """Derive a key for encrypting the DEK in the license."""
        from cryptography.hazmat.primitives.kdf.hkdf import HKDF
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=license_id.encode(),
            info=b"vllm-model-security-dek-key",
            backend=default_backend(),
        )
        return hkdf.derive(vendor_secret)

    @staticmethod
    def encrypt_dek_for_license(dek: bytes, dek_key: bytes) -> str:
        """Encrypt DEK using derived key."""
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        import os
        aesgcm = AESGCM(dek_key)
        nonce = os.urandom(12)
        encrypted = aesgcm.encrypt(nonce, dek, None)
        return base64.b64encode(nonce + encrypted).decode()

    @staticmethod
    def decrypt_dek_from_license(dek_encrypted_b64: str, dek_key: bytes) -> bytes:
        """Decrypt DEK from license."""
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        data = base64.b64decode(dek_encrypted_b64)
        nonce = data[:12]
        ciphertext = data[12:]
        aesgcm = AESGCM(dek_key)
        return aesgcm.decrypt(nonce, ciphertext, None)
