"""Content scanner — detect leaked secrets and PII in tool responses.

Part of the Obol AI Agent Immune System SDK. Scans PostToolUse responses
and tool inputs for accidentally leaked secrets and personally identifiable
information.

Runs entirely on the user's device. Zero external dependencies (stdlib only).
Never stores or transmits the actual content — only detection flag names.

Detection layers:
  - Secret patterns: API keys, private keys, database URLs, tokens
  - PII patterns: SSN, credit cards, bulk emails, phone numbers, internal IPs
  - Input scanning: secrets being written to files via Write/Edit tools

Usage:
    from obol.hooks.content_scanner import scan_response, scan_input

    # In PostToolUse handler:
    result = scan_response("Bash", tool_response)
    # {"secrets_found": ["aws_access_key"], "pii_found": [], ...}

    result = scan_input("Write", tool_input)
    # {"secrets_found": ["private_key_rsa"], "pii_found": [], ...}
"""
import re


# ---- Secret patterns --------------------------------------------------------
# Each tuple: (flag_name, compiled_regex)
# Flag names are returned in results; the matched content is NEVER included.

SECRET_PATTERNS = [
    # Cloud provider keys
    ("aws_access_key", re.compile(r'AKIA[A-Z0-9]{16}')),
    ("aws_secret_key", re.compile(
        r'(?:aws_secret_access_key|AWS_SECRET_ACCESS_KEY)\s*[=:]\s*[A-Za-z0-9/+=]{40}'
    )),
    ("gcp_service_account", re.compile(r'"type"\s*:\s*"service_account"')),
    ("azure_connection_string", re.compile(
        r'(?:AccountKey|SharedAccessKey)\s*=\s*[A-Za-z0-9+/=]{44,}'
    )),

    # API keys (common formats)
    ("anthropic_api_key", re.compile(r'sk-ant-[a-zA-Z0-9_-]{20,}')),
    ("openai_api_key", re.compile(r'sk-(?!ant-|_(?:live|test)_)[a-zA-Z0-9]{20,}')),
    ("stripe_key", re.compile(r'(?:sk|pk)_(?:live|test)_[a-zA-Z0-9]{24,}')),
    ("github_token", re.compile(r'(?:ghp|gho|ghs|ghr)_[a-zA-Z0-9]{36,}')),
    ("github_fine_grained", re.compile(r'github_pat_[a-zA-Z0-9_]{22,}')),
    ("slack_token", re.compile(r'xox[bpors]-[a-zA-Z0-9-]{10,}')),
    ("slack_webhook", re.compile(
        r'https://hooks\.slack\.com/services/T[a-zA-Z0-9_]+/B[a-zA-Z0-9_]+/[a-zA-Z0-9_]+'
    )),
    ("discord_webhook", re.compile(
        r'https://discord(?:app)?\.com/api/webhooks/\d+/[a-zA-Z0-9_-]+'
    )),
    ("discord_token", re.compile(r'[MN][A-Za-z\d]{23,}\.[\w-]{6}\.[\w-]{27,}')),
    ("telegram_bot_token", re.compile(r'\d{8,10}:[A-Za-z0-9_-]{35}')),

    # Additional cloud/SaaS provider tokens
    ("huggingface_token", re.compile(r'hf_[a-zA-Z0-9]{34,}')),
    ("vercel_token", re.compile(r'vercel_[a-zA-Z0-9]{24,}')),
    ("supabase_key", re.compile(r'(?:sbp|sbn)_[a-zA-Z0-9]{40,}')),
    ("sendgrid_key", re.compile(r'SG\.[a-zA-Z0-9_-]{22,}\.[a-zA-Z0-9_-]{43,}')),
    ("twilio_key", re.compile(r'SK[a-f0-9]{32}')),
    ("cloudflare_token", re.compile(r'(?:cf_|cloudflare_)[a-zA-Z0-9_-]{37,}')),
    ("digitalocean_token", re.compile(r'dop_v1_[a-f0-9]{64}')),
    ("npm_token", re.compile(r'npm_[a-zA-Z0-9]{36,}')),

    # Obol API keys
    ("obol_api_key", re.compile(r'ob_(?:live|test)_[a-f0-9]{32,}')),

    # HashiCorp Vault tokens (hvs. for service tokens, s. for legacy)
    ("vault_token", re.compile(r'(?:hvs\.[a-zA-Z0-9_-]{24,}|s\.[a-zA-Z0-9]{24,})')),
    # 1Password service account tokens
    ("onepassword_sa_token", re.compile(r'ops_eyJ[a-zA-Z0-9_-]{50,}')),
    # Doppler tokens (dp.st.dev.xxx, dp.ct.xxx, etc.)
    ("doppler_token", re.compile(r'dp\.(?:st|ct|sa|scim|audit)\.[a-zA-Z0-9_.]{10,}')),
    # Netlify personal access tokens
    ("netlify_token", re.compile(r'nfp_[a-zA-Z0-9]{40,}')),
    # Infisical tokens
    ("infisical_token", re.compile(r'st\.[a-zA-Z0-9]{20,}\.[a-zA-Z0-9]{6,}')),

    # Additional provider tokens
    ("pypi_token", re.compile(r'pypi-[A-Za-z0-9_-]{50,}')),
    ("gitlab_token", re.compile(r'glpat-[a-zA-Z0-9_-]{20,}')),
    ("databricks_token", re.compile(r'dapi[a-f0-9]{32,}')),
    ("linear_api_key", re.compile(r'lin_api_[a-zA-Z0-9]{40,}')),
    ("resend_api_key", re.compile(r're_[a-zA-Z0-9]{20,}_[a-zA-Z0-9]+')),
    ("google_api_key", re.compile(r'AIza[a-zA-Z0-9_-]{35}')),

    # Private keys
    ("private_key_rsa", re.compile(r'-----BEGIN (?:RSA )?PRIVATE KEY-----')),
    ("private_key_ec", re.compile(r'-----BEGIN EC PRIVATE KEY-----')),
    ("private_key_openssh", re.compile(r'-----BEGIN OPENSSH PRIVATE KEY-----')),
    ("private_key_pgp", re.compile(r'-----BEGIN PGP PRIVATE KEY BLOCK-----')),

    # Database connection strings
    ("database_url", re.compile(
        r'(?:postgres(?:ql)?|mysql|mongodb(?:\+srv)?|redis)://[^\s"\']+:[^\s"\']+@[^\s"\']+'
    )),

    # JWT tokens (in responses, indicates token leakage)
    ("jwt_token", re.compile(
        r'eyJ[a-zA-Z0-9_-]{10,}\.eyJ[a-zA-Z0-9_-]{10,}\.[a-zA-Z0-9_-]{10,}'
    )),

    # Generic secrets in key=value format
    ("generic_api_key", re.compile(
        r'(?:api[_-]?key|apikey)\s*[=:\uff1d\uff1a]\s*["\']?[a-zA-Z0-9_-]{20,}["\']?',
        re.IGNORECASE,
    )),
    ("generic_secret", re.compile(
        r'(?:secret)\s*[=:\uff1d\uff1a]\s*["\']?[a-zA-Z0-9_+/=-]{20,}["\']?',
        re.IGNORECASE,
    )),
    ("generic_password", re.compile(
        r'(?:password|passwd)\s*[=:\uff1d\uff1a]\s*["\']?[^\s"\']{8,}["\']?',
        re.IGNORECASE,
    )),
]


# ---- Base64 entropy detection -----------------------------------------------
# Detects base64-encoded secrets in tool responses.
# Tries to decode and re-scan, catching encoded API keys and private keys.

_BASE64_BLOB_RE = re.compile(r'[A-Za-z0-9+/]{40,}={0,2}')
_HEX_BLOB_RE = re.compile(r'(?:0x)?([0-9a-fA-F]{40,})')
# Match strings with 3+ URL-encoded chars mixed with regular text (20+ total chars)
_URL_ENCODED_RE = re.compile(r'(?:[A-Za-z0-9_.~-]|%[0-9a-fA-F]{2}){20,}')

def _check_base64_secrets(text):
    """Try to decode base64 blobs in text and re-scan for secrets.

    Returns a list of secret flag names found in decoded content.
    Only checks blobs of 40+ chars to reduce false positives.
    """
    import base64
    found = []
    blobs = _BASE64_BLOB_RE.findall(text)
    # Check up to 20 blobs to prevent multi-chunk secret splitting evasion
    # Cap individual blob size at 10KB to prevent memory amplification DoS
    for blob in blobs[:20]:
        if len(blob) > 14000:  # ~10KB decoded (base64 expands ~33%)
            continue
        try:
            raw = base64.b64decode(blob, validate=True)
            decoded = raw.decode("utf-8", errors="replace")
            if len(decoded) < 10:
                continue
            # Re-scan decoded content with secret patterns
            matched = False
            for flag_name, pattern in SECRET_PATTERNS:
                if pattern.search(decoded):
                    found.append(f"base64_encoded_{flag_name}")
                    matched = True
                    break
            if not matched:
                # Try gzip decompression (base64(gzip(secret)) evasion)
                try:
                    import zlib
                    decompressed = zlib.decompress(raw, 16 + zlib.MAX_WBITS)
                    decoded_gz = decompressed[:10000].decode("utf-8", errors="replace")
                    for flag_name, pattern in SECRET_PATTERNS:
                        if pattern.search(decoded_gz):
                            found.append(f"base64_gzip_{flag_name}")
                            matched = True
                            break
                except Exception:
                    pass
            if not matched:
                # Flag high-entropy binary blobs that can't be pattern-matched
                # (>10% replacement chars suggests binary secret masking)
                replacement_frac = decoded.count("\ufffd") / len(decoded) if decoded else 0
                if replacement_frac > 0.10 and len(raw) >= 32:
                    found.append("suspicious_binary_blob")
        except Exception:
            continue
    return found


def _check_hex_secrets(text):
    """Try to decode hex-encoded blobs and re-scan for secrets.

    Returns a list of secret flag names found in decoded content.
    Catches: echo "736b2d616e742d..." (hex-encoded API keys).
    """
    found = []
    blobs = _HEX_BLOB_RE.findall(text)
    for blob in blobs[:10]:
        if len(blob) > 2048:
            continue
        try:
            import binascii
            raw = binascii.unhexlify(blob)
            decoded = raw.decode("utf-8", errors="replace")
            if len(decoded) < 10:
                continue
            # Only flag if the decoded text looks like ASCII (not random binary)
            printable_frac = sum(1 for c in decoded if c.isprintable()) / len(decoded)
            if printable_frac < 0.7:
                continue
            for flag_name, pattern in SECRET_PATTERNS:
                if pattern.search(decoded):
                    found.append(f"hex_encoded_{flag_name}")
                    break
        except Exception:
            continue
    return found


def _check_url_encoded_secrets(text):
    """Try to URL-decode encoded content and re-scan for secrets.

    Returns a list of secret flag names found in decoded content.
    Catches: sk%2Dant%2D1234567890abcdef
    """
    found = []
    blobs = _URL_ENCODED_RE.findall(text)
    for blob in blobs[:10]:
        if len(blob) > 2048:
            continue
        # Require at least 3 percent-encoded chars to avoid false positives
        pct_count = blob.count('%')
        if pct_count < 3:
            continue
        try:
            from urllib.parse import unquote
            decoded = unquote(blob)
            if decoded == blob or len(decoded) < 10:
                continue
            for flag_name, pattern in SECRET_PATTERNS:
                if pattern.search(decoded):
                    found.append(f"url_encoded_{flag_name}")
                    break
        except Exception:
            continue
    return found


# ---- PII patterns -----------------------------------------------------------
# SSN and credit card use regex; email/phone/IP use custom logic with thresholds.

_SSN_RE = re.compile(r'\b\d{3}-\d{2}-\d{4}\b')

_CREDIT_CARD_RE = re.compile(
    r'\b(?:4\d{3}|5[1-5]\d{2}|3[47]\d{2}|6(?:011|5\d{2}))[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b'
)

# Email regex with bounded local-part length to prevent ReDoS on long alpha strings.
# RFC 5321 limits local-part to 64 chars; we use 128 for safety.
_EMAIL_RE = re.compile(r'[a-zA-Z0-9._%+-]{1,128}@[a-zA-Z0-9.-]{1,255}\.[a-zA-Z]{2,}')

_PHONE_RE = re.compile(r'\b(?:\+1[- ]?)?\(?\d{3}\)?[- ]?\d{3}[- ]?\d{4}\b')

_PRIVATE_IP_RE = re.compile(
    r'\b(?:10\.\d{1,3}\.\d{1,3}\.\d{1,3}'
    r'|172\.(?:1[6-9]|2\d|3[01])\.\d{1,3}\.\d{1,3}'
    r'|192\.168\.\d{1,3}\.\d{1,3})\b'
)

# Thresholds for bulk PII detection (single occurrences are normal)
_EMAIL_THRESHOLD = 5
_PHONE_THRESHOLD = 3
_PRIVATE_IP_THRESHOLD = 3


# ---- Tools that produce scannable text content ------------------------------
# Only these tools return text worth scanning for leaked secrets/PII.
_RESPONSE_TOOLS = {"Bash", "Read", "WebFetch", "WebSearch", "Grep"}

# Tools where input content should be scanned (agent writing secrets to files)
_INPUT_TOOLS = {"Write", "Edit", "NotebookEdit"}


# ---- Luhn algorithm ---------------------------------------------------------

def _luhn_check(number_str):
    """Validate a credit card number with the Luhn algorithm.

    Returns True if the number passes the checksum, False otherwise.
    Strips non-digit characters (spaces, dashes) before validation.
    """
    digits = [int(d) for d in number_str if d.isdigit()]
    if len(digits) < 13 or len(digits) > 19:
        return False
    checksum = 0
    for i, d in enumerate(reversed(digits)):
        if i % 2 == 1:
            d *= 2
            if d > 9:
                d -= 9
        checksum += d
    return checksum % 10 == 0


# ---- Text extraction --------------------------------------------------------

def _extract_text(tool_name, tool_response, max_bytes):
    """Extract scannable text from a tool response.

    Handles dict responses (stdout, content, text, output fields) and
    plain strings. Returns the text truncated to max_bytes, and a bool
    indicating whether truncation occurred.
    """
    text = ""
    if isinstance(tool_response, dict):
        # Bash: stdout + stderr
        if tool_name == "Bash":
            text = (tool_response.get("stdout", "") or "") + "\n" + (
                tool_response.get("stderr", "") or ""
            )
        else:
            # Try common response fields in order of likelihood
            for field in ("content", "text", "output", "stdout"):
                val = tool_response.get(field)
                if val and isinstance(val, str):
                    text = val
                    break
            # If nothing found, skip (don't serialize the whole dict)
    elif isinstance(tool_response, str):
        text = tool_response

    if not text:
        return "", False

    truncated = len(text) > max_bytes
    if truncated:
        # Scan head + tail to catch secrets hidden after the scan limit
        tail_size = min(5000, max_bytes // 10)
        combined = text[:max_bytes] + "\n" + text[-tail_size:]
        return combined, True
    return text, False


def _extract_input_text(tool_name, tool_input):
    """Extract scannable text from a tool input dict.

    Only extracts the content fields that represent file data being written.
    Returns the extracted text (may be empty).
    """
    if not tool_input or not isinstance(tool_input, dict):
        return ""

    if tool_name == "Write":
        return tool_input.get("content", "") or ""
    if tool_name == "Edit":
        return tool_input.get("new_string", "") or ""
    if tool_name == "NotebookEdit":
        return tool_input.get("new_source", "") or ""
    return ""


# ---- Scanning logic ---------------------------------------------------------

def _scan_text(text):
    """Scan a text string for secrets and PII.

    Returns a dict:
        secrets_found: list of unique secret flag names
        pii_found: list of unique PII flag names
        secret_count: total number of individual secret matches
        pii_count: total number of individual PII matches
    """
    secrets_found = []
    pii_found = []
    secret_count = 0
    pii_count = 0

    # Secret scanning — use search() not findall() for performance.
    # findall() scans entire text even after first match; on 50KB inputs
    # with 30+ patterns this is a bottleneck. We only need boolean detection.
    seen_secret_types = set()
    for flag_name, pattern in SECRET_PATTERNS:
        if pattern.search(text):
            secret_count += 1
            if flag_name not in seen_secret_types:
                seen_secret_types.add(flag_name)
                secrets_found.append(flag_name)

    # PII scanning: SSN (always flag)
    ssn_matches = _SSN_RE.findall(text)
    if ssn_matches:
        pii_found.append("ssn")
        pii_count += len(ssn_matches)

    # PII scanning: credit cards (Luhn-validated)
    cc_matches = _CREDIT_CARD_RE.findall(text)
    if cc_matches:
        valid_count = sum(1 for m in cc_matches if _luhn_check(m))
        if valid_count > 0:
            pii_found.append("credit_card")
            pii_count += valid_count

    # PII scanning: emails (only flag if bulk, >threshold unique)
    email_matches = _EMAIL_RE.findall(text)
    if email_matches:
        unique_emails = set(e.lower() for e in email_matches)
        if len(unique_emails) > _EMAIL_THRESHOLD:
            pii_found.append("email_addresses")
            pii_count += len(unique_emails)

    # PII scanning: phone numbers (only flag if >threshold found)
    phone_matches = _PHONE_RE.findall(text)
    if len(phone_matches) > _PHONE_THRESHOLD:
        pii_found.append("phone_number")
        pii_count += len(phone_matches)

    # PII scanning: internal IPs (only flag if >threshold unique)
    ip_matches = _PRIVATE_IP_RE.findall(text)
    if ip_matches:
        unique_ips = set(ip_matches)
        if len(unique_ips) > _PRIVATE_IP_THRESHOLD:
            pii_found.append("internal_ip_addresses")
            pii_count += len(unique_ips)

    # Base64 entropy detection: decode and re-scan for encoded secrets
    try:
        b64_secrets = _check_base64_secrets(text)
        for flag in b64_secrets:
            if flag not in seen_secret_types:
                seen_secret_types.add(flag)
                secrets_found.append(flag)
                secret_count += 1
    except Exception:
        pass

    # Hex-encoded secret detection
    try:
        hex_secrets = _check_hex_secrets(text)
        for flag in hex_secrets:
            if flag not in seen_secret_types:
                seen_secret_types.add(flag)
                secrets_found.append(flag)
                secret_count += 1
    except Exception:
        pass

    # URL-encoded secret detection
    try:
        url_secrets = _check_url_encoded_secrets(text)
        for flag in url_secrets:
            if flag not in seen_secret_types:
                seen_secret_types.add(flag)
                secrets_found.append(flag)
                secret_count += 1
    except Exception:
        pass

    return {
        "secrets_found": secrets_found,
        "pii_found": pii_found,
        "secret_count": secret_count,
        "pii_count": pii_count,
    }


# ---- Public API -------------------------------------------------------------

def scan_response(tool_name, tool_response, max_scan_bytes=50_000):
    """Scan a tool response for leaked secrets and PII.

    Called from the PostToolUse handler after a tool runs. Inspects the
    response text for accidentally exposed secrets and personally
    identifiable information.

    Args:
        tool_name: The name of the tool that produced the response
            (e.g. "Bash", "Read", "WebFetch", "Grep").
        tool_response: The tool's response — either a dict with content
            fields or a plain string.
        max_scan_bytes: Maximum bytes to scan (default 50KB). Responses
            larger than this are truncated before scanning.

    Returns:
        dict with keys:
            secrets_found: list of unique secret type names detected
            pii_found: list of unique PII type names detected
            secret_count: total number of individual secret matches
            pii_count: total number of individual PII matches
            scan_truncated: True if response exceeded max_scan_bytes

    Never includes the actual secret or PII content. Empty lists mean
    nothing was found. Never raises exceptions.
    """
    empty = {
        "secrets_found": [],
        "pii_found": [],
        "secret_count": 0,
        "pii_count": 0,
        "scan_truncated": False,
    }

    try:
        # Only scan tools that produce text content
        if tool_name not in _RESPONSE_TOOLS:
            return empty

        text, truncated = _extract_text(tool_name, tool_response, max_scan_bytes)
        if not text:
            return empty

        result = _scan_text(text)
        result["scan_truncated"] = truncated
        return result

    except Exception:
        return empty


def scan_input(tool_name, tool_input, max_scan_bytes=50_000):
    """Scan tool input for secrets being written to files.

    Called from the PostToolUse handler to check whether the agent is
    writing secrets into files (e.g. committing API keys to source code).
    Only inspects Write, Edit, and NotebookEdit tool inputs.

    Args:
        tool_name: The name of the tool (e.g. "Write", "Edit",
            "NotebookEdit").
        tool_input: The tool's input dict containing the content
            being written.
        max_scan_bytes: Maximum bytes to scan (default 50KB).

    Returns:
        dict with same structure as scan_response(). Never includes
        the actual secret or PII content. Never raises exceptions.
    """
    empty = {
        "secrets_found": [],
        "pii_found": [],
        "secret_count": 0,
        "pii_count": 0,
        "scan_truncated": False,
    }

    try:
        # Only scan write-type tools
        if tool_name not in _INPUT_TOOLS:
            return empty

        text = _extract_input_text(tool_name, tool_input)
        if not text:
            return empty

        truncated = len(text) > max_scan_bytes
        if truncated:
            # Scan head + tail to catch secrets at end of large files
            tail_size = min(5000, max_scan_bytes // 10)
            text = text[:max_scan_bytes] + "\n" + text[-tail_size:]
        # (short text passes through unchanged)

        result = _scan_text(text)
        result["scan_truncated"] = truncated
        return result

    except Exception:
        return empty
