"""Canarytoken management for Warden.

A canarytoken is a plausible-looking credential file (fake AWS access key,
fake SSH private key, fake .env) planted in a location an AI agent is
likely to touch during reconnaissance. Each file contains a unique marker
string that Warden's policy engine detects on read.

Unlike the `secret-access` rule which flags reads of sensitive-named files,
canarytokens trigger even when the file name looks mundane — the agent
has to actually OPEN the file to hit the marker.

Registered canaries live in `$HOME/.prismor/canaries.json`. Each entry:
  { "id": "<uuid>", "path": "<file>", "type": "aws|ssh|env|generic",
    "marker": "<random-string>", "webhook": "<optional-url>",
    "created": "<iso-timestamp>" }

On read, Warden raises a CRITICAL finding and (if a webhook is configured)
POSTs a tamper-evident JSON payload to it.
"""
from __future__ import annotations

import json
import os
import secrets
import string
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

CANARY_REGISTRY = Path.home() / ".prismor" / "canaries.json"
MARKER_PREFIX = "PRISMOR-CANARY-"


# ── Templates ──────────────────────────────────────────────────────────────

_AWS_TEMPLATE = """[default]
# AWS credentials — DO NOT SHARE
aws_access_key_id = AKIA{key_suffix}
aws_secret_access_key = {secret_body}
# {marker}
"""

_SSH_TEMPLATE = """-----BEGIN OPENSSH PRIVATE KEY-----
b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW
{body}
{marker}
-----END OPENSSH PRIVATE KEY-----
"""

_ENV_TEMPLATE = """# Application secrets — DO NOT COMMIT
DATABASE_URL=postgres://app:{db_pw}@db.internal.example.com:5432/prod
STRIPE_SECRET_KEY=sk_live_{stripe_suffix}
JWT_SIGNING_KEY={jwt}
# {marker}
"""

_GENERIC_TEMPLATE = """# {marker}
# Generated by Prismor Warden — if you read this through an AI agent,
# alert your security team.
"""


def _random_suffix(n: int = 32, alphabet: str = string.ascii_uppercase + string.digits) -> str:
    return "".join(secrets.choice(alphabet) for _ in range(n))


def _new_marker() -> str:
    return f"{MARKER_PREFIX}{secrets.token_hex(12).upper()}"


# ── Registry ───────────────────────────────────────────────────────────────

def _load_registry() -> List[Dict[str, Any]]:
    if not CANARY_REGISTRY.exists():
        return []
    try:
        return json.loads(CANARY_REGISTRY.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return []


def _save_registry(entries: List[Dict[str, Any]]) -> None:
    CANARY_REGISTRY.parent.mkdir(parents=True, exist_ok=True)
    CANARY_REGISTRY.write_text(json.dumps(entries, indent=2), encoding="utf-8")
    # Registry contains markers that bypass detection if leaked — restrict.
    try:
        os.chmod(CANARY_REGISTRY, 0o600)
    except OSError:
        pass


def list_canaries() -> List[Dict[str, Any]]:
    """Return all registered canaries (with marker values redacted)."""
    entries = _load_registry()
    return [
        {**e, "marker": e["marker"][: len(MARKER_PREFIX) + 4] + "…"}
        for e in entries
    ]


def get_markers() -> List[str]:
    """Return every active canary marker — used by the policy engine."""
    return [e["marker"] for e in _load_registry()]


# ── Plant / unplant ────────────────────────────────────────────────────────

def _render(template_type: str, marker: str) -> str:
    if template_type == "aws":
        return _AWS_TEMPLATE.format(
            key_suffix=_random_suffix(16),
            secret_body=_random_suffix(40, string.ascii_letters + string.digits + "/+"),
            marker=marker,
        )
    if template_type == "ssh":
        # 12 random-looking base64 lines. Length doesn't matter for realism,
        # but it needs to be enough to pass a cursory glance.
        body = "\n".join(_random_suffix(70, string.ascii_letters + string.digits + "+/") for _ in range(12))
        return _SSH_TEMPLATE.format(body=body, marker=marker)
    if template_type == "env":
        return _ENV_TEMPLATE.format(
            db_pw=_random_suffix(24),
            stripe_suffix=_random_suffix(24, string.ascii_letters + string.digits),
            jwt=_random_suffix(44, string.ascii_letters + string.digits + "-_"),
            marker=marker,
        )
    return _GENERIC_TEMPLATE.format(marker=marker)


def plant(
    path: Path,
    template: str = "generic",
    webhook: Optional[str] = None,
    force: bool = False,
) -> Dict[str, Any]:
    """Plant a canarytoken at ``path`` with the given template type.

    Returns the new registry entry. Raises FileExistsError if the path
    already exists and ``force=False``.
    """
    if template not in {"aws", "ssh", "env", "generic"}:
        raise ValueError(f"unknown template: {template}")

    path = path.expanduser().resolve()
    if path.exists() and not force:
        raise FileExistsError(f"{path} already exists (use --force to overwrite)")

    marker = _new_marker()
    content = _render(template, marker)

    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")
    try:
        os.chmod(path, 0o600 if template in {"aws", "ssh", "env"} else 0o644)
    except OSError:
        pass

    entry = {
        "id": secrets.token_hex(8),
        "path": str(path),
        "type": template,
        "marker": marker,
        "webhook": webhook,
        "created": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    }
    entries = _load_registry()
    entries.append(entry)
    _save_registry(entries)
    return entry


def unplant(identifier: str) -> Optional[Dict[str, Any]]:
    """Remove a canary by id or path. Returns the removed entry or None."""
    entries = _load_registry()
    kept: List[Dict[str, Any]] = []
    removed: Optional[Dict[str, Any]] = None
    ident = str(Path(identifier).expanduser().resolve()) if "/" in identifier else identifier
    for e in entries:
        if e["id"] == identifier or e["path"] == ident:
            removed = e
            try:
                Path(e["path"]).unlink(missing_ok=True)
            except OSError:
                pass
            continue
        kept.append(e)
    _save_registry(kept)
    return removed


# ── Detection hook (used by policy_engine) ─────────────────────────────────

def check_content_for_markers(text: str) -> Optional[str]:
    """Return the first canary marker found in ``text``, or None.

    Used by PostToolUse content scanning: if a file read returned a canary
    marker in stdout, the agent definitely touched a honey-token.
    """
    if not text or MARKER_PREFIX not in text:
        return None
    for marker in get_markers():
        if marker in text:
            return marker
    return None


def check_path_is_canary(path: str) -> Optional[Dict[str, Any]]:
    """Return the canary registry entry if ``path`` matches, else None.

    Used by PreToolUse file_read hooks: if the agent is about to read a
    registered canary path, we can trigger before any content is returned.
    """
    if not path:
        return None
    try:
        resolved = str(Path(path).expanduser().resolve())
    except (OSError, ValueError):
        resolved = path
    for e in _load_registry():
        if e["path"] == resolved or e["path"] == path:
            return e
    return None


# ── Webhook beacon ─────────────────────────────────────────────────────────

def beacon(entry: Dict[str, Any], event: str, context: Optional[Dict[str, Any]] = None) -> bool:
    """POST a JSON payload to the entry's webhook. Returns True on 2xx."""
    url = entry.get("webhook")
    if not url:
        return False
    try:
        import urllib.request
        import urllib.error
        payload = {
            "event": event,
            "canary_id": entry["id"],
            "canary_path": entry["path"],
            "canary_type": entry["type"],
            "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "hostname": os.uname().nodename if hasattr(os, "uname") else "unknown",
            "context": context or {},
        }
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            url, data=data, headers={"Content-Type": "application/json"}, method="POST"
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            return 200 <= resp.status < 300
    except Exception:
        return False
