Simplify Python proxy automation

strix/tools/proxy/caido_api.py

@@ -0,0 +1,437 @@
+"""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",
+]