app/security.py

import secrets
import hashlib
import hmac
from typing import Optional

from app.config import Settings


def generate_session_id() -> str:
    return "sess_" + secrets.token_hex(16)


def generate_worker_id() -> str:
    return "worker_" + secrets.token_hex(12)


def generate_job_id() -> str:
    return "job_" + secrets.token_hex(16)


def generate_receipt_id() -> str:
    return "rcpt_" + secrets.token_hex(16)


def generate_session_secret() -> str:
    return secrets.token_urlsafe(32)


def hash_payload(payload_bytes: bytes) -> str:
    return hashlib.sha256(payload_bytes).hexdigest()


def hash_json(payload_object: dict) -> str:
    import json
    canonical = json.dumps(payload_object, sort_keys=True, separators=(',', ':'))
    return hashlib.sha256(canonical.encode('utf-8')).hexdigest()


def constant_time_compare(a: str, b: str) -> bool:
    return hmac.compare_digest(a.encode(), b.encode())


def create_worker_join_token(session_id: str) -> str:
    raw = f"{session_id}:{secrets.token_urlsafe(16)}"
    return hashlib.sha256(raw.encode()).hexdigest()[:32]


def verify_worker_join_token(session_id: str, token: str) -> bool:
    # In production, store and check against a database of valid tokens
    # For v1, accept any well-formed token and rely on session_id matching
    return len(token) == 32 and token.isalnum()


def validate_payload_size(payload_bytes: bytes) -> bool:
    return len(payload_bytes) <= Settings.get_max_payload_bytes()


def validate_allowed_job_type(job_type: str) -> bool:
    from app.models import JobType
    try:
        JobType(job_type)
        return True
    except ValueError:
        return False


def validate_origin(origin: Optional[str]) -> bool:
    if not origin:
        return True
    # In production, check against allowed origins
    allowed = {"https://huggingface.co", "https://spaces.huggingface.tech"}
    return any(origin.startswith(a) for a in allowed) or origin.startswith("http://localhost")


def sanitize_public_error(error: str) -> str:
    # Do not leak internal details
    sensitive = ["secret", "token", "password", "key", "private"]
    low = error.lower()
    if any(s in low for s in sensitive):
        return "Internal error"
    return error[:200]


# Receipt signing helpers
_server_nonce_cache: dict = {}


def generate_server_nonce() -> str:
    nonce = secrets.token_hex(16)
    _server_nonce_cache[nonce] = True
    return nonce


def verify_device_signature(public_key: str, message: str, signature: str) -> bool:
    # v1: Ed25519 verification if keys present; otherwise trust-on-first-use
    if not public_key or not signature:
        return not Settings.is_signature_required()
    try:
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
        from cryptography.exceptions import InvalidSignature
        pk_bytes = bytes.fromhex(public_key)
        pk = Ed25519PublicKey.from_public_bytes(pk_bytes)
        sig_bytes = bytes.fromhex(signature)
        pk.verify(sig_bytes, message.encode('utf-8'))
        return True
    except Exception:
        return not Settings.is_signature_required()


def canonicalize_receipt_for_signing(receipt: dict) -> str:
    import json
    # Exclude signatures from canonicalization
    filtered = {k: v for k, v in receipt.items() if k not in ("device_signature", "server_signature", "receipt_hash")}
    return json.dumps(filtered, sort_keys=True, separators=(',', ':'))