strix/strix/tools/proxy/caido_api.py

"""Shared Caido proxy helpers and sandbox-importable ``caido_api`` module."""

from __future__ import annotations

import asyncio
import json
import os
import time
import urllib.request
from typing import TYPE_CHECKING, Any, Literal
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

from caido_sdk_client import Client, TokenAuthOptions
from caido_sdk_client.types import (
    ConnectionInfoInput,
    CreateReplaySessionFromRaw,
    CreateReplaySessionOptions,
    CreateScopeOptions,
    ReplaySendOptions,
    RequestGetOptions,
    UpdateScopeOptions,
)


if TYPE_CHECKING:
    from caido_sdk_client import Client as CaidoClient


RequestPart = Literal["request", "response"]
SortBy = Literal[
    "timestamp",
    "host",
    "method",
    "path",
    "status_code",
    "response_time",
    "response_size",
    "source",
]
SortOrder = Literal["asc", "desc"]
ScopeAction = Literal["get", "list", "create", "update", "delete"]

_DEFAULT_CAIDO_URL = "http://127.0.0.1:48080"
_CLIENT_CACHE: dict[str, Client] = {}
_REQ_FIELD_MAP: dict[SortBy, tuple[str, str]] = {
    "timestamp": ("req", "created_at"),
    "host": ("req", "host"),
    "method": ("req", "method"),
    "path": ("req", "path"),
    "source": ("req", "source"),
    "status_code": ("resp", "code"),
    "response_time": ("resp", "roundtrip"),
    "response_size": ("resp", "length"),
}


def caido_url() -> str:
    """Return the in-sandbox Caido endpoint used by ``caido_api``."""
    return os.environ.get("STRIX_CAIDO_URL", _DEFAULT_CAIDO_URL).rstrip("/")


def _graphql_url() -> str:
    base_url = caido_url()
    parsed = urlparse(base_url)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        raise ValueError(f"Invalid Caido URL: {base_url}")
    return f"{base_url}/graphql"


def _login_as_guest() -> str:
    body = json.dumps({"query": "mutation { loginAsGuest { token { accessToken } } }"}).encode(
        "utf-8"
    )
    req = urllib.request.Request(  # noqa: S310
        _graphql_url(),
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=10) as resp:  # noqa: S310  # nosec B310
        payload = json.loads(resp.read())
    return str(payload["data"]["loginAsGuest"]["token"]["accessToken"])


async def get_client() -> Client:
    """Return a connected Caido SDK client for the local sandbox sidecar."""
    if client := _CLIENT_CACHE.get("default"):
        return client

    token = await asyncio.to_thread(_login_as_guest)
    client = Client(caido_url(), auth=TokenAuthOptions(token=token))
    await client.connect()
    _CLIENT_CACHE["default"] = client
    return client


async def close_client() -> None:
    """Close the cached sandbox Caido client, if one was opened."""
    client = _CLIENT_CACHE.pop("default", None)
    if client is None:
        return
    await client.aclose()


async def list_requests_with_client(
    client: CaidoClient,
    *,
    httpql_filter: str | None = None,
    first: int = 50,
    after: str | None = None,
    sort_by: SortBy = "timestamp",
    sort_order: SortOrder = "desc",
    scope_id: str | None = None,
) -> Any:
    builder = client.request.list().first(first)
    if httpql_filter:
        builder = builder.filter(httpql_filter)
    if after:
        builder = builder.after(after)
    if scope_id:
        builder = builder.scope(scope_id)
    target, field = _REQ_FIELD_MAP[sort_by]
    builder = (builder.descending if sort_order == "desc" else builder.ascending)(target, field)
    return await builder.execute()


async def get_request_with_client(
    client: CaidoClient,
    request_id: str,
    *,
    part: RequestPart = "request",
) -> Any:
    opts = RequestGetOptions(
        request_raw=(part == "request"),
        response_raw=(part == "response"),
    )
    return await client.request.get(request_id, opts)


def build_raw_request(
    *,
    method: str,
    url: str,
    headers: dict[str, str],
    body: str,
) -> tuple[ConnectionInfoInput, bytes]:
    parsed = urlparse(url)
    if not parsed.scheme or not parsed.netloc:
        raise ValueError(f"Invalid URL: {url}")
    is_tls = parsed.scheme.lower() == "https"
    host = parsed.hostname or ""
    port = parsed.port or (443 if is_tls else 80)
    path = parsed.path or "/"
    if parsed.query:
        path = f"{path}?{parsed.query}"

    final_headers = {**headers}
    final_headers.setdefault("Host", parsed.netloc)
    final_headers.setdefault("User-Agent", "strix")
    if body and "Content-Length" not in {k.title() for k in final_headers}:
        final_headers["Content-Length"] = str(len(body.encode("utf-8")))

    lines = [f"{method.upper()} {path} HTTP/1.1"]
    lines.extend(f"{k}: {v}" for k, v in final_headers.items())
    raw = ("\r\n".join(lines) + "\r\n\r\n" + body).encode("utf-8")
    return ConnectionInfoInput(host=host, port=port, is_tls=is_tls), raw


def parse_raw_request(raw_content: str) -> dict[str, Any]:
    lines = raw_content.split("\n")
    request_line = lines[0].strip().split(" ")
    if len(request_line) < 2:
        raise ValueError("Invalid request line format")
    method, url_path = request_line[0], request_line[1]

    parsed_headers: dict[str, str] = {}
    body_start = 0
    for i, line in enumerate(lines[1:], 1):
        if line.strip() == "":
            body_start = i + 1
            break
        if ":" in line:
            key, value = line.split(":", 1)
            parsed_headers[key.strip()] = value.strip()

    body = "\n".join(lines[body_start:]).strip() if body_start < len(lines) else ""
    return {"method": method, "url_path": url_path, "headers": parsed_headers, "body": body}


def full_url_from_components(
    original: Any,
    components: dict[str, Any],
    modifications: dict[str, Any],
) -> str:
    if "url" in modifications:
        return str(modifications["url"])
    headers = components["headers"]
    host_header = headers.get("Host") or original.host
    scheme = "https" if original.is_tls else "http"
    return f"{scheme}://{host_header}{components['url_path']}"


def apply_modifications(
    components: dict[str, Any],
    modifications: dict[str, Any],
    full_url: str,
) -> dict[str, Any]:
    headers = dict(components["headers"])
    body = components["body"]
    final_url = full_url

    if "params" in modifications:
        parsed = urlparse(final_url)
        existing = {k: v[0] if v else "" for k, v in parse_qs(parsed.query).items()}
        existing.update(modifications["params"])
        final_url = urlunparse(parsed._replace(query=urlencode(existing)))
    if "headers" in modifications:
        headers.update(modifications["headers"])
    if "body" in modifications:
        body = modifications["body"]
    if "cookies" in modifications:
        cookies: dict[str, str] = {}
        if headers.get("Cookie"):
            for cookie in headers["Cookie"].split(";"):
                if "=" in cookie:
                    k, v = cookie.split("=", 1)
                    cookies[k.strip()] = v.strip()
        cookies.update(modifications["cookies"])
        headers["Cookie"] = "; ".join(f"{k}={v}" for k, v in cookies.items())

    return {
        "method": components["method"],
        "url": final_url,
        "headers": headers,
        "body": body,
    }


async def replay_send_raw(
    client: CaidoClient,
    *,
    raw: bytes,
    connection: ConnectionInfoInput,
) -> dict[str, Any]:
    started = time.time()
    session = await client.replay.sessions.create(
        CreateReplaySessionOptions(
            request_source=CreateReplaySessionFromRaw(raw=raw, connection=connection),
        ),
    )
    result = await client.replay.send(
        session.id,
        ReplaySendOptions(raw=raw, connection=connection),
    )
    elapsed_ms = int((time.time() - started) * 1000)
    response_raw = result.entry.response_raw if hasattr(result.entry, "response_raw") else None
    return {
        "session_id": str(session.id),
        "status": result.status,
        "error": result.error,
        "elapsed_ms": elapsed_ms,
        "response_raw": response_raw,
    }


async def scope_list(client: CaidoClient) -> Any:
    return await client.scope.list()


async def scope_get(client: CaidoClient, scope_id: str) -> Any:
    return await client.scope.get(scope_id)


async def scope_create(
    client: CaidoClient,
    *,
    name: str,
    allowlist: list[str] | None = None,
    denylist: list[str] | None = None,
) -> Any:
    return await client.scope.create(
        CreateScopeOptions(
            name=name,
            allowlist=list(allowlist or []),
            denylist=list(denylist or []),
        ),
    )


async def scope_update(
    client: CaidoClient,
    scope_id: str,
    *,
    name: str,
    allowlist: list[str] | None = None,
    denylist: list[str] | None = None,
) -> Any:
    return await client.scope.update(
        scope_id,
        UpdateScopeOptions(
            name=name,
            allowlist=list(allowlist or []),
            denylist=list(denylist or []),
        ),
    )


async def scope_delete(client: CaidoClient, scope_id: str) -> None:
    await client.scope.delete(scope_id)


async def list_requests(
    *,
    httpql_filter: str | None = None,
    first: int = 50,
    after: str | None = None,
    sort_by: SortBy = "timestamp",
    sort_order: SortOrder = "desc",
    scope_id: str | None = None,
) -> Any:
    """List captured HTTP requests from sandbox Python."""
    return await list_requests_with_client(
        await get_client(),
        httpql_filter=httpql_filter,
        first=first,
        after=after,
        sort_by=sort_by,
        sort_order=sort_order,
        scope_id=scope_id,
    )


async def view_request(request_id: str, *, part: RequestPart = "request") -> Any:
    """Return one captured request/response from sandbox Python."""
    return await get_request_with_client(await get_client(), request_id, part=part)


async def send_request(
    method: str,
    url: str,
    *,
    headers: dict[str, str] | None = None,
    body: str = "",
) -> dict[str, Any]:
    """Send an arbitrary raw HTTP request through Caido Replay."""
    connection, raw = build_raw_request(
        method=method,
        url=url,
        headers=headers or {},
        body=body,
    )
    return await replay_send_raw(await get_client(), raw=raw, connection=connection)


async def repeat_request(
    request_id: str,
    *,
    modifications: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Replay a captured request after applying request modifications."""
    mods = modifications or {}
    result = await get_request_with_client(await get_client(), request_id, part="request")
    if result is None or result.request.raw is None:
        raise ValueError(f"Request {request_id} not found")

    original = result.request
    raw_str = result.request.raw.decode("utf-8", errors="replace")
    components = parse_raw_request(raw_str)
    full_url = full_url_from_components(original, components, mods)
    modified = apply_modifications(components, mods, full_url)
    connection, raw = build_raw_request(
        method=modified["method"],
        url=modified["url"],
        headers=modified["headers"],
        body=modified["body"],
    )
    return await replay_send_raw(await get_client(), raw=raw, connection=connection)


async def scope_rules(
    action: ScopeAction,
    *,
    allowlist: list[str] | None = None,
    denylist: list[str] | None = None,
    scope_id: str | None = None,
    scope_name: str | None = None,
) -> Any:
    """Manage Caido scope rules from sandbox Python."""
    client = await get_client()
    if action == "list":
        result = await scope_list(client)
    elif action == "get":
        if not scope_id:
            raise ValueError("scope_id required for get")
        result = await scope_get(client, scope_id)
    elif action == "create":
        if not scope_name:
            raise ValueError("scope_name required for create")
        result = await scope_create(
            client,
            name=scope_name,
            allowlist=allowlist,
            denylist=denylist,
        )
    elif action == "update":
        if not scope_id or not scope_name:
            raise ValueError("scope_id and scope_name required for update")
        result = await scope_update(
            client,
            scope_id,
            name=scope_name,
            allowlist=allowlist,
            denylist=denylist,
        )
    elif action == "delete":
        if not scope_id:
            raise ValueError("scope_id required for delete")
        await scope_delete(client, scope_id)
        result = {"deleted": scope_id}
    else:
        raise ValueError(f"Unknown action: {action}")
    return result


__all__ = [
    "RequestPart",
    "ScopeAction",
    "SortBy",
    "SortOrder",
    "close_client",
    "get_client",
    "list_requests",
    "repeat_request",
    "scope_rules",
    "send_request",
    "view_request",
]