"""Process-wide ``socket.connect`` audit hook for defense-in-depth
egress control.

Why this exists
---------------
Every explicit PEP we wrote gates a single known caller — ContentFetcher,
evaluate_llm_endpoint, MCP download_content, and so on. They are the
primary line of defense and they remain the primary line of defense.
This module is the secondary line — the "what about the code paths we
haven't reviewed yet?" net. It fires on EVERY outbound socket.connect
inside the process, so a langchain tool registered by an MCP server, a
new contributor reaching for ``requests.get`` directly, or a third-party
library that opens its own connection all get caught by the same rule
the user configured.

Threat model
------------
- Catches: forgotten PEP coverage, new code paths that don't know about
  the policy, third-party library egress, prompt-injection that steers a
  tool into raw HTTP, ``urllib.request.urlopen``, ``httpx``, ``aiohttp``,
  ``socket.socket().connect()`` — anything that ultimately calls
  ``socket.connect`` (almost everything network-bound in CPython).
- Does NOT catch: a determined adversary with code execution in the LDR
  process. They can ``clear_active_context()`` from inside the process,
  monkey-patch the hook, ``addaudithook`` a passthrough, or open a file
  descriptor out-of-band. For a real boundary, layer OS-level controls
  (network namespaces, firewall rules, restricted Docker) — see
  SECURITY.md.

Activation model
----------------
- ``install_audit_hook()`` is INSTALLED ONCE at app/library startup
  (idempotent). PEP 578 audit hooks cannot be removed — that is the
  design, not a limitation. Once it is in, it stays.
- The hook is INACTIVE by default: when no EgressContext is registered
  for the current thread, every connect passes through unmodified. This
  is the right default — random scripts importing the package and
  pytest collectors that touch a socket would otherwise raise
  ``PolicyDeniedError`` mid-test.
- The hook becomes ACTIVE for a thread when a worker calls
  ``set_active_context(ctx)``. Clear with ``clear_active_context()`` in
  a finally to prevent leak across thread-pool reuse. The
  ``database.thread_local_session._ThreadCleanup`` exit handler already
  does this for the standard research worker lifecycle.
"""

from __future__ import annotations

import socket
import sys
import threading
from contextlib import contextmanager
from typing import Optional

from loguru import logger


# ---------------------------------------------------------------------------
# Thread-local state
# ---------------------------------------------------------------------------

# Per-thread "currently active" EgressContext. The audit hook reads this; a
# worker (research run, news scheduler, library indexer, etc.) sets it at
# task entry and clears it in finally.
_thread_local = threading.local()

# Re-entrancy guard. The hook itself may indirectly trigger a socket call
# (e.g., if a future logger ever ships a record through a network sink)
# and without this we would recurse on every connect.
_hook_reentry = threading.local()


def set_active_context(ctx) -> None:
    """Register ``ctx`` as the active EgressContext for the current
    thread. Subsequent ``socket.connect`` calls from this thread are
    gated through the policy until ``clear_active_context()`` is called.

    Idempotent: setting None is equivalent to clearing.
    """
    if ctx is None:
        clear_active_context()
        return
    # An ADAPTIVE scope must be resolved to a concrete scope BEFORE it is
    # stored: the audit hook only gates PRIVATE_ONLY/STRICT, so an ADAPTIVE
    # context reaching the hook would make the backstop a silent no-op.
    # Production always builds contexts via context_from_snapshot (which
    # resolves ADAPTIVE); this fail-fast guard stops a future direct
    # construction or test from disarming the net unnoticed.
    try:
        from .policy import EgressScope

        if getattr(ctx, "scope", None) == EgressScope.ADAPTIVE:
            raise ValueError(
                "EgressContext.scope must be resolved to a concrete scope "
                "before activation; got ADAPTIVE. Build the context via "
                "context_from_snapshot(), which resolves ADAPTIVE."
            )
    except ImportError:  # pragma: no cover - policy import always available
        pass
    _thread_local.egress_context = ctx


def clear_active_context() -> None:
    """Clear the active EgressContext for the current thread. Safe to
    call when no context is set.
    """
    # ``hasattr`` + ``del`` rather than ``setattr(None)`` so the
    # ``getattr(..., None)`` branch in the hook short-circuits without
    # touching the policy module at all.
    if hasattr(_thread_local, "egress_context"):
        del _thread_local.egress_context


def get_active_context():
    """Return the active EgressContext for this thread, or ``None``."""
    return getattr(_thread_local, "egress_context", None)


@contextmanager
def active_egress_context(ctx):
    """Context manager: set the active EgressContext for the duration of
    the ``with`` block and clear it on exit (even on exception).

    Useful for unit tests, ad-hoc scripts, and any worker that wants
    automatic cleanup without writing a try/finally.
    """
    # Save/restore the previous context rather than unconditionally clearing,
    # so nesting an inner ``active_egress_context`` doesn't wipe an outer
    # one's context on exit (the single thread-local slot otherwise loses the
    # parent — a nested chat-in-research or test would leave the thread with
    # NO active context, silently disarming the backstop for the parent's
    # remaining work).
    previous = get_active_context()
    set_active_context(ctx)
    try:
        yield
    finally:
        if previous is not None:
            set_active_context(previous)
        else:
            clear_active_context()


# ---------------------------------------------------------------------------
# The audit hook
# ---------------------------------------------------------------------------

_INSTALLED = False
_INSTALL_LOCK = threading.Lock()


def install_audit_hook() -> None:
    """Install the ``sys.audit`` hook. Idempotent — calling multiple
    times is a no-op after the first. Per PEP 578 audit hooks cannot be
    uninstalled, so this is install-once-and-keep-forever; the activation
    gate is the per-thread context, not the hook itself.
    """
    global _INSTALLED
    with _INSTALL_LOCK:
        if _INSTALLED:
            return
        sys.addaudithook(_audit_hook)
        _INSTALLED = True


def is_installed() -> bool:
    return _INSTALLED


def _extract_host(address) -> Optional[str]:
    """Extract the destination host from a ``socket.connect`` address
    tuple. Returns ``None`` for shapes we don't recognise (caller passes
    through — failing to extract is never a reason to block).

    AF_INET addresses are ``(host, port)``. AF_INET6 addresses are
    ``(host, port, flowinfo, scope_id)``. ``host`` is a string in both.
    """
    if not isinstance(address, tuple) or not address:
        return None
    host = address[0]
    # CPython accepts a bytes/bytearray host in ``socket.connect`` and fires
    # the ``socket.connect`` audit event with the raw bytes (verified). Decode
    # so a bytes-encoded IP/host literal is classified instead of silently
    # passing through the hook — otherwise a library that connects with a
    # bytes host bypasses the PRIVATE_ONLY/STRICT backstop entirely.
    if isinstance(host, (bytes, bytearray)):
        try:
            host = bytes(host).decode("ascii")
        except (UnicodeDecodeError, ValueError):
            return None
    if not isinstance(host, str):
        return None
    return host


def _audit_hook(event: str, args: tuple) -> None:
    """sys.audit hook. Only watches ``socket.connect``; ignores every
    other event. Raises ``PolicyDeniedError`` when the active context
    refuses the destination.

    Failure mode: this hook deliberately does NOT swallow exceptions
    raised by the policy lookup. If the policy module itself errors,
    the connect fails loudly so the operator notices — the
    catch-and-continue alternative would degrade the secondary line of
    defense into "silently disabled when there's a bug." That is the
    opposite of what a defense-in-depth net is for.
    """
    if event != "socket.connect":
        return

    # No active policy → fast path. This is the common case for any
    # thread that is not running a research worker (e.g., the Flask
    # request handler, sqlalchemy connection pools, the news scheduler
    # in between jobs). Skipping here keeps the hook overhead at
    # roughly two attribute lookups per connect.
    ctx = getattr(_thread_local, "egress_context", None)
    if ctx is None:
        return

    # Re-entrancy guard. Anything that runs inside the hook (logger
    # binding, exception construction, attribute access) might
    # indirectly trigger another connect (e.g. via a future remote log
    # sink). Without this, the hook would recurse and stack-overflow.
    if getattr(_hook_reentry, "active", False):
        return

    # Arm the re-entrancy guard BEFORE touching ``args``/``sock``/``address``
    # or importing the policy module: any of those attribute accesses could
    # (under instrumentation or a future remote log sink) trigger another
    # socket.connect, and the guard must already be set to break the
    # recursion. The try/finally guarantees the guard is cleared on EVERY
    # exit path — the early returns below and the PolicyDeniedError raise
    # included — so a blocked connect doesn't wedge the hook off for the
    # rest of the thread's life.
    _hook_reentry.active = True
    try:
        sock, address = args

        # Only gate network sockets. AF_UNIX, AF_NETLINK, AF_PACKET, etc.
        # are off-network and should not pay the policy cost.
        family = getattr(sock, "family", None)
        if family not in (socket.AF_INET, socket.AF_INET6):
            return

        host = _extract_host(address)
        if not host:
            return

        # Scope-restriction: this hook is a "no-public-egress" net, NOT a
        # general re-implementation of every PEP at the socket layer.
        # PRIVATE_ONLY and STRICT are the two scopes that explicitly say
        # "no public host should be reached." PUBLIC_ONLY governs search-
        # ENGINE selection (the user wants public sources only) and is NOT
        # supposed to block legitimate local infrastructure traffic — local
        # Ollama, local embeddings, the user's own settings DB. Same logic
        # for BOTH. Applying evaluate_url under those scopes would refuse
        # the local LLM connect with ``scope_mismatch_public_only``, which
        # is a false positive against the user's actual intent.
        from .policy import EgressScope

        if ctx.scope not in (EgressScope.PRIVATE_ONLY, EgressScope.STRICT):
            return

        # Imported lazily to avoid a hard dependency at hook-install
        # time. The hook is installed early during app startup and the
        # egress_policy module pulls in network_utils, ip_ranges, etc.
        # — none of which are problematic, but the lazy import keeps
        # the install path minimal.
        from .policy import (
            PolicyDeniedError,
            evaluate_url,
        )

        # Synthesize a URL so evaluate_url's existing host classification
        # logic does the work (private-IP check, NAT64, local_hostnames
        # match, DNS resolution if the host is not already an IP). The
        # scheme is fixed at http:// because evaluate_url only inspects
        # the host part — the value here is irrelevant. IPv6 literals
        # MUST be bracketed (``[::1]``) or urlsplit's hostname parser
        # returns None and the policy then refuses on ``no_hostname``.
        url_host = f"[{host}]" if family == socket.AF_INET6 else host
        decision = evaluate_url(f"http://{url_host}", ctx)
        if not decision.allowed:
            logger.bind(policy_audit=True).warning(
                "audit hook blocked socket.connect",
                host=host,
                reason=decision.reason,
                scope=ctx.scope.value,
            )
            raise PolicyDeniedError(decision, target=host)
    finally:
        _hook_reentry.active = False
