#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
// PURPOSE: Cryptographic Proof-of-Origin for PULSE Evidence.
// [Phase 1.2: Cryptographic Integrity - HARDENED]
"""

import hmac
import hashlib
import json
import os
import sys
import time
from pathlib import Path
from typing import Dict, Any, Optional
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.exceptions import InvalidSignature

# Add Utilities to path

class PulseCrypto:
    """
    // Handles signing and verification of evidence items.
    // HARDENED: No hardcoded secrets allowed.
    // KEY VAULTING: Secret is wiped from environment after load.
    """

    # M2: Anti-replay window (5 minutes). Signatures older than this are rejected.
    REPLAY_WINDOW = 300

    def __init__(self, secret_key: Optional[str] = None, wipe_env: bool = False):
        # 1. Load Symmetric Key (HMAC) for Agent-to-Agent trust
        self.secret = secret_key or os.environ.get("PULSE_BRAIN_KEY")
        if not self.secret:
            raise RuntimeError(
                "PULSE_BRAIN_KEY not found in environment. "
                "Set it in .env or pass secret_key to PulseCrypto()."
            )

        # 2. Key Vaulting: Optionally wipe from environment to prevent process leaks.
        #    Disabled by default because CLI tools need subprocesses to inherit the key.
        if wipe_env and "PULSE_BRAIN_KEY" in os.environ:
            del os.environ["PULSE_BRAIN_KEY"]

        self.secret_bytes = self.secret.encode()

        # M1: Key version for rotation support
        self.key_version = int(os.environ.get("PULSE_KEY_VERSION", "1"))

    def sign(self, item: Dict[str, Any]) -> str:
        """Generates a HMAC-SHA256 signature for a dictionary.
        M1: Includes key version. M2: Includes timestamp for anti-replay."""
        clean_item = {k: v for k, v in item.items()
                      if k not in ("_signature", "_hash", "_key_version", "_signed_at")}
        data = json.dumps(clean_item, sort_keys=True).encode()
        return hmac.new(self.secret_bytes, data, hashlib.sha256).hexdigest()

    def verify(self, item: Dict[str, Any]) -> bool:
        """Verifies the HMAC signature of an evidence item.
        M1: Checks key version compatibility. M2: Rejects stale signatures."""
        signature = item.get("_signature")
        if not signature:
            return False

        # M1: Key version check — reject items signed with a different key version
        item_key_version = item.get("_key_version", 1)
        if item_key_version != self.key_version:
            return False

        # M2: Anti-replay — reject signatures older than REPLAY_WINDOW
        signed_at = item.get("_signed_at")
        if signed_at is not None:
            try:
                age = time.time() - float(signed_at)
                if age > self.REPLAY_WINDOW:
                    return False
            except (ValueError, TypeError):
                return False

        expected = self.sign(item)
        return hmac.compare_digest(signature, expected)

    def seal(self, item: Dict[str, Any], agent_id: str) -> Dict[str, Any]:
        """Seals an item with agent ID, key version, timestamp, and signature.
        M1: Embeds key version. M2: Embeds signature timestamp."""
        item["_agent_id"] = agent_id
        item["_key_version"] = self.key_version
        item["_signed_at"] = time.time()
        item["_hash"] = hashlib.sha256(json.dumps(item, sort_keys=True).encode()).hexdigest()
        item["_signature"] = self.sign(item)
        return item

    # --- SENSORY SEALING (Phase 3.3) ---

    def get_agent_key(self, agent_id: str) -> bytes:
        """// Derives an agent-specific HMAC key from the master secret."""
        return hmac.new(self.secret_bytes, agent_id.encode(), hashlib.sha256).digest()

    def sign_string(self, text: str, agent_id: str) -> str:
        """// Signs a plain string (e.g., a log line) using an agent-specific key."""
        agent_key = self.get_agent_key(agent_id)
        return hmac.new(agent_key, text.encode(), hashlib.sha256).hexdigest()

    def verify_string(self, text: str, signature: str, agent_id: str) -> bool:
        """// Verifies a plain string against its signature and agent ID."""
        expected = self.sign_string(text, agent_id)
        return hmac.compare_digest(signature, expected)

    # --- NEW: RSA VERIFICATION FOR HUMAN AUTHORITY ---
    
    @staticmethod
    def verify_human_signature(message: str, signature_hex: str, public_key_pem: str) -> bool:
        """
        // Verifies a Founder's RSA signature against the whitelist.
        // Used for locking/unlocking critical proposals.
        """
        try:
            public_key = serialization.load_pem_public_key(
                public_key_pem.encode()
            )
            
            signature_bytes = bytes.fromhex(signature_hex)
            message_bytes = message.encode()

            public_key.verify(
                signature_bytes,
                message_bytes,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
            return True
        except (InvalidSignature, ValueError):
            return False
        except Exception as e:
            print(f"Crypto Error: {e}")
            return False

    @staticmethod
    def sign_with_rsa(data: bytes, private_key_path: Path) -> bytes:
        """
        // Signs data using an RSA Private Key (PEM).
        // Used by ComplianceValidator to issue Proof of Compliance.
        """
        with open(private_key_path, "rb") as key_file:
            private_key = serialization.load_pem_private_key(
                key_file.read(),
                password=None
            )

        signature = private_key.sign(
            data,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        return signature

# --- MODULE LEVEL WRAPPERS (For easier imports) ---

def sign_data(data: bytes, private_key_path: Path) -> bytes:
    """// Module-level wrapper for RSA signing."""
    return PulseCrypto.sign_with_rsa(data, private_key_path)

def verify_poc(data: bytes, signature_hex: str, public_key_pem: str) -> bool:
    """// Module-level wrapper for RSA verification."""
    return PulseCrypto.verify_human_signature(data.decode(), signature_hex, public_key_pem)


def sign_item(item: dict, agent_id: str) -> dict:
    """HMAC-seal a dict item for integrity tracking. No-op if PULSE_BRAIN_KEY is not set."""
    if not isinstance(item, dict) or not agent_id:
        return item
    try:
        crypto = PulseCrypto()
        return crypto.seal(item, agent_id)
    except (RuntimeError, ImportError):
        return item  # Graceful degradation — key not set

if __name__ == "__main__":
    # Test Hardening
    print("Testing PulseCrypto Hardening...")
    try:
        # This should FAIL if env var is missing
        crypto = PulseCrypto() 
        print("FAIL: PulseCrypto initialized without key!")
    except SystemExit:
        print("PASS: PulseCrypto refused to start without key.")