"""
crypto_attestation.py — Cryptographic attestation for SichGate findings.

Signs the full findings report and individual status-transition events with
ed25519, creating an immutable chain of custody that prevents post-hoc
tampering.  The public key is embedded in every signed artefact so any
third-party tool can verify the signature without contacting SichGate.

Chain of custody:
  1. Findings report signed at scan completion (``sign_findings_report``).
  2. Each status transition (NEW→MITIGATED, etc.) signed individually
     (``sign_status_transition``).
  3. ``verify_report`` / ``verify_external`` allow offline verification.

Key management:
  - ``CryptoAttestation()`` generates a fresh ed25519 key pair.
  - ``CryptoAttestation(key_path=…)`` loads an existing PEM private key.
  - Call ``save_key(path)`` to persist a generated key for future runs.
"""
from __future__ import annotations

import base64
import hashlib
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
    load_pem_private_key,
)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _canonical_json(obj: dict) -> bytes:
    """
    Stable, deterministic JSON serialisation for signing.

    ``sort_keys=True`` + no whitespace ensures the same bytes for the same
    logical content regardless of dict insertion order.
    """
    return json.dumps(obj, sort_keys=True, separators=(",", ":"),
                      ensure_ascii=True).encode("utf-8")


def _b64(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("ascii")


def _sha256_hex(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


# ---------------------------------------------------------------------------
# CryptoAttestation
# ---------------------------------------------------------------------------

class CryptoAttestation:
    """
    ed25519 signing authority for SichGate compliance artefacts.

    Example::

        att = CryptoAttestation()
        att.save_key("sichgate_key.pem")          # persist for subsequent runs

        signed = att.sign_findings_report(findings_dict)
        assert att.verify_report(signed)           # True for self-signed

        # Hand signed["attestation"]["signature"] + public_key_b64 to auditor.
        assert CryptoAttestation.verify_external(signed)
    """

    def __init__(self, key_path: Optional[str] = None) -> None:
        if key_path and Path(key_path).exists():
            self._private_key: Ed25519PrivateKey = self._load_key(key_path)
        else:
            self._private_key = Ed25519PrivateKey.generate()
        self._public_key: Ed25519PublicKey = self._private_key.public_key()

    # ── Key management ───────────────────────────────────────────────────────

    @staticmethod
    def _load_key(path: str) -> Ed25519PrivateKey:
        pem = Path(path).read_bytes()
        return load_pem_private_key(pem, password=None)  # type: ignore[return-value]

    def save_key(self, path: str) -> None:
        """Write the private key as PEM to *path*.  Keep this secret."""
        pem = self._private_key.private_bytes(
            encoding=Encoding.PEM,
            format=PrivateFormat.PKCS8,
            encryption_algorithm=NoEncryption(),
        )
        Path(path).write_bytes(pem)

    def public_key_pem(self) -> str:
        """Return the public key in PEM format (safe to share with auditors)."""
        return self._public_key.public_bytes(
            encoding=Encoding.PEM,
            format=PublicFormat.SubjectPublicKeyInfo,
        ).decode("ascii")

    def public_key_b64(self) -> str:
        """Return the raw public key as URL-safe base64 (compact JSON form)."""
        raw = self._public_key.public_bytes(
            encoding=Encoding.Raw,
            format=PublicFormat.Raw,
        )
        return _b64(raw)

    # ── Signing ──────────────────────────────────────────────────────────────

    def _sign_bytes(self, data: bytes) -> str:
        return _b64(self._private_key.sign(data))

    # ── Public API ───────────────────────────────────────────────────────────

    def sign_findings_report(self, findings: dict) -> dict:
        """
        Sign a complete findings report dict.

        Returns a new dict containing the original findings plus an
        ``attestation`` sub-object with:

        - ``attestation_id``   — unique ID for this attestation event
        - ``timestamp``        — ISO-8601 UTC timestamp
        - ``signature``        — ed25519 signature (URL-safe base64)
        - ``public_key_b64``   — verifying public key (base64 raw)
        - ``public_key_pem``   — verifying public key (PEM)
        - ``payload_hash``     — SHA-256 of the canonical signed payload
        - ``algorithm``        — "ed25519"
        """
        attest_id = str(uuid.uuid4())
        timestamp = _now_iso()

        payload_to_sign = {
            "findings": findings,
            "attestation_id": attest_id,
            "timestamp": timestamp,
        }
        canonical = _canonical_json(payload_to_sign)
        sig = self._sign_bytes(canonical)
        payload_hash = _sha256_hex(canonical)

        return {
            **findings,
            "attestation": {
                "attestation_id": attest_id,
                "timestamp": timestamp,
                "signature": sig,
                "public_key_b64": self.public_key_b64(),
                "public_key_pem": self.public_key_pem(),
                "payload_hash": payload_hash,
                "algorithm": "ed25519",
            },
        }

    def sign_status_transition(
        self,
        finding_id: str,
        from_status: str,
        new_status: str,
        actor: Optional[str] = None,
    ) -> dict:
        """
        Sign an audit trail status transition event.

        Returns a signed dict suitable for appending to the
        ``status_history`` list of an audit trail.
        """
        transition_id = str(uuid.uuid4())
        timestamp = _now_iso()

        transition = {
            "transition_id": transition_id,
            "finding_id": finding_id,
            "from_status": from_status,
            "to_status": new_status,
            "timestamp": timestamp,
            "actor": actor or "sichgate",
        }
        canonical = _canonical_json(transition)
        sig = self._sign_bytes(canonical)

        return {
            **transition,
            "signature": sig,
            "public_key_b64": self.public_key_b64(),
            "algorithm": "ed25519",
        }

    def build_audit_trail(
        self,
        finding_id: str,
        initial_status: str = "NEW",
        actor: Optional[str] = None,
    ) -> dict:
        """
        Create a new audit trail with a signed creation event.

        Returns an audit trail dict with ``finding_id``, ``created_at``,
        ``current_status``, and ``status_history``.
        """
        creation_entry = self.sign_status_transition(
            finding_id=finding_id,
            from_status="",
            new_status=initial_status,
            actor=actor,
        )
        return {
            "finding_id": finding_id,
            "created_at": creation_entry["timestamp"],
            "current_status": initial_status,
            "status_history": [creation_entry],
        }

    def append_transition(
        self,
        audit_trail: dict,
        new_status: str,
        actor: Optional[str] = None,
    ) -> dict:
        """
        Append a signed status transition to an existing audit trail.

        Returns the updated audit trail dict (original is not mutated).
        """
        current = audit_trail.get("current_status", "NEW")
        entry = self.sign_status_transition(
            finding_id=audit_trail["finding_id"],
            from_status=current,
            new_status=new_status,
            actor=actor,
        )
        return {
            **audit_trail,
            "current_status": new_status,
            "status_history": audit_trail["status_history"] + [entry],
        }

    # ── Verification ─────────────────────────────────────────────────────────

    def verify_report(self, signed_report: dict) -> bool:
        """
        Verify the signature of a report signed by *this* instance.

        Returns True if signature is valid and payload hash matches.
        """
        return self._verify_with_public_key(signed_report, self._public_key)

    @staticmethod
    def verify_external(signed_report: dict) -> bool:
        """
        Verify a report using the public key embedded in the report itself.

        This method can be called by any party who has the signed JSON —
        no private key or SichGate installation required.

        Returns True if signature and payload hash are both valid.
        """
        att = signed_report.get("attestation", {})
        pub_b64 = att.get("public_key_b64")
        if not pub_b64:
            return False
        # Pad base64 to multiple of 4
        raw_pub = base64.urlsafe_b64decode(pub_b64 + "==")
        pub_key = Ed25519PublicKey.from_public_bytes(raw_pub)
        return CryptoAttestation._verify_with_public_key(signed_report, pub_key)

    @staticmethod
    def _verify_with_public_key(signed_report: dict, pub_key: Ed25519PublicKey) -> bool:
        att = signed_report.get("attestation", {})
        sig_b64 = att.get("signature")
        attest_id = att.get("attestation_id")
        timestamp = att.get("timestamp")
        payload_hash = att.get("payload_hash")

        if not all([sig_b64, attest_id, timestamp]):
            return False

        findings = {k: v for k, v in signed_report.items() if k != "attestation"}
        payload_to_verify = {
            "findings": findings,
            "attestation_id": attest_id,
            "timestamp": timestamp,
        }
        canonical = _canonical_json(payload_to_verify)

        # Verify hash first (fast fail)
        if payload_hash and _sha256_hex(canonical) != payload_hash:
            return False

        sig_bytes = base64.urlsafe_b64decode(sig_b64 + "==")
        try:
            pub_key.verify(sig_bytes, canonical)
            return True
        except Exception:
            return False
