mirror of
https://github.com/usestrix/strix.git
synced 2026-08-20 10:33:34 +02:00
498 lines
20 KiB
Python
498 lines
20 KiB
Python
"""UI-independent state and command controller for interactive Strix clients."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import math
|
|
import webbrowser
|
|
from collections.abc import Awaitable, Callable
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from strix.config import load_settings
|
|
from strix.config.models import is_recommended_or_frontier_model
|
|
from strix.config.settings import DEFAULT_MAX_TURNS
|
|
from strix.interface.tui.backend.live_view import TuiLiveView
|
|
from strix.interface.tui.backend.projection import (
|
|
MAX_TERMINAL_EVENTS,
|
|
MAX_TERMINAL_VULNERABILITIES,
|
|
SCAN_MODES,
|
|
SCOPE_MODES,
|
|
bounded_state_projection,
|
|
collection_item_projection,
|
|
sanitize_terminal_text,
|
|
terminal_projection,
|
|
)
|
|
from strix.interface.utils import is_subscription_run
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
import argparse
|
|
|
|
from strix.report.state import ReportState
|
|
|
|
|
|
_STOPPABLE_AGENT_STATUSES = frozenset({"running", "waiting", "budget_paused"})
|
|
|
|
ChangeCallback = Callable[[], None]
|
|
StartCallback = Callable[[bool], Awaitable[None]]
|
|
QuitCallback = Callable[[], Awaitable[None]]
|
|
|
|
|
|
class TuiController:
|
|
"""Own setup state and expose serializable scan state to any TUI."""
|
|
|
|
def __init__(
|
|
self,
|
|
args: argparse.Namespace,
|
|
*,
|
|
live_view: TuiLiveView | None = None,
|
|
coordinator: Any = None,
|
|
report_state: ReportState | None = None,
|
|
on_start: StartCallback | None = None,
|
|
on_quit: QuitCallback | None = None,
|
|
on_change: ChangeCallback | None = None,
|
|
) -> None:
|
|
self.args = args
|
|
self.live_view = live_view or TuiLiveView()
|
|
self.coordinator = coordinator
|
|
self.report_state = report_state
|
|
self.scan_loop: asyncio.AbstractEventLoop | None = None
|
|
self.setup_mode = bool(args.needs_setup)
|
|
self.scan_started = not self.setup_mode
|
|
self._start_in_progress = False
|
|
self.scan_state = "setup" if self.setup_mode else "running"
|
|
self.targets = [
|
|
str(target["original"])
|
|
for target in args.targets_info
|
|
if isinstance(target, dict) and target.get("original")
|
|
]
|
|
instruction = args.instruction
|
|
self.instruction = instruction.strip() if isinstance(instruction, str) else ""
|
|
requested_scan_mode = str(args.scan_mode)
|
|
self.scan_mode = requested_scan_mode if requested_scan_mode in SCAN_MODES else "deep"
|
|
raw_budget = args.max_budget_usd
|
|
self.max_budget_usd = (
|
|
float(raw_budget)
|
|
if isinstance(raw_budget, int | float)
|
|
and not isinstance(raw_budget, bool)
|
|
and math.isfinite(float(raw_budget))
|
|
and raw_budget > 0
|
|
else None
|
|
)
|
|
raw_turns = args.max_turns
|
|
self.max_turns = (
|
|
raw_turns
|
|
if isinstance(raw_turns, int) and not isinstance(raw_turns, bool) and raw_turns > 0
|
|
else DEFAULT_MAX_TURNS
|
|
)
|
|
requested_scope = str(args.scope_mode)
|
|
self.scope_mode = requested_scope if requested_scope in SCOPE_MODES else "auto"
|
|
raw_diff_base = args.diff_base
|
|
self.diff_base = raw_diff_base.strip() if isinstance(raw_diff_base, str) else None
|
|
# Host directory mounted for the agent to work in when the scan has no
|
|
# target, set only once the user confirms it. It is a workspace, not a
|
|
# target: it carries no scan scope, and the instruction is the only
|
|
# source of truth for what to do.
|
|
self.workspace_mount: str | None = None
|
|
# A target-less launch enters the live view and asks there before
|
|
# anything is prepared; this holds the directory awaiting that answer.
|
|
self.pending_workspace_mount: str | None = None
|
|
self._pending_verify = True
|
|
self.messages: list[dict[str, str]] = []
|
|
self._next_message_id = 1
|
|
self.error: str | None = None
|
|
self.viewer_status = "idle"
|
|
self.viewer_url: str | None = None
|
|
self._viewer_httpd: Any = None
|
|
self._on_start = on_start
|
|
self._on_quit = on_quit
|
|
self._on_change = on_change
|
|
|
|
def set_change_callback(self, callback: ChangeCallback) -> None:
|
|
self._on_change = callback
|
|
|
|
def notify_changed(self) -> None:
|
|
if self._on_change is not None:
|
|
self._on_change()
|
|
|
|
def set_runtime(
|
|
self,
|
|
*,
|
|
report_state: ReportState | None = None,
|
|
scan_loop: asyncio.AbstractEventLoop | None = None,
|
|
) -> None:
|
|
if report_state is not None:
|
|
self.report_state = report_state
|
|
if scan_loop is not None:
|
|
self.scan_loop = scan_loop
|
|
|
|
def begin_preparation(self) -> None:
|
|
"""Mark a directly-launched run as preparing behind the live TUI."""
|
|
self.scan_state = "preparing"
|
|
self.notify_changed()
|
|
|
|
def fail_preparation(self, detail: str) -> None:
|
|
self.scan_state = "failed"
|
|
self.error = detail
|
|
self.notify_changed()
|
|
|
|
def enter_setup(self) -> None:
|
|
"""Return a session to the start screen, e.g. on a declined mount."""
|
|
self.setup_mode = True
|
|
self.scan_started = False
|
|
self.scan_state = "setup"
|
|
self.notify_changed()
|
|
|
|
def add_message(self, text: str, level: str = "info") -> None:
|
|
self._append_message(text, level)
|
|
self.notify_changed()
|
|
|
|
def _append_message(self, text: str, level: str) -> None:
|
|
self.messages.append(
|
|
{
|
|
"id": f"message-{self._next_message_id}",
|
|
"text": sanitize_terminal_text(text),
|
|
"level": sanitize_terminal_text(level),
|
|
}
|
|
)
|
|
self._next_message_id += 1
|
|
self.messages = self.messages[-200:]
|
|
|
|
def snapshot(self) -> dict[str, Any]:
|
|
"""Return small mutable state; histories are streamed as collections."""
|
|
model = ""
|
|
with contextlib.suppress(Exception):
|
|
model = (load_settings().llm.model or "").strip()
|
|
usage: dict[str, Any] = {}
|
|
if self.report_state is not None:
|
|
usage = dict(self.report_state.get_total_llm_usage())
|
|
subscription = False
|
|
with contextlib.suppress(Exception):
|
|
subscription = is_subscription_run(self.report_state)
|
|
model_warning = ""
|
|
if model and not is_recommended_or_frontier_model(model):
|
|
model_warning = (
|
|
f"{model} is not a recommended frontier model; pentest quality could be degraded"
|
|
)
|
|
state = {
|
|
"setup_mode": self.setup_mode,
|
|
"scan_started": self.scan_started,
|
|
"scan_state": self.scan_state,
|
|
"targets": [
|
|
terminal_projection(target, max_string=128) for target in self.targets[:16]
|
|
],
|
|
"target_count": len(self.targets),
|
|
"working_dir": str(Path.cwd()),
|
|
"pending_mount": self.pending_workspace_mount or "",
|
|
"instruction": terminal_projection(self.instruction, max_string=2 * 1024),
|
|
"scan_mode": self.scan_mode,
|
|
"max_budget_usd": self.max_budget_usd,
|
|
"max_turns": self.max_turns,
|
|
"scope_mode": self.scope_mode,
|
|
"diff_base": terminal_projection(self.diff_base, max_string=256),
|
|
"model": terminal_projection(model, max_string=256),
|
|
"model_warning": terminal_projection(model_warning, max_string=512),
|
|
"caido_url": terminal_projection(
|
|
getattr(self.report_state, "caido_url", None), max_string=1024
|
|
),
|
|
"messages": [
|
|
{
|
|
"id": str(message.get("id", ""))[:64],
|
|
"text": terminal_projection(message.get("text", ""), max_string=256),
|
|
"level": str(message.get("level", "info"))[:32],
|
|
}
|
|
for message in self.messages[-10:]
|
|
],
|
|
"usage": terminal_projection(usage, max_string=256, max_items=20),
|
|
"subscription": subscription,
|
|
"viewer_status": self.viewer_status,
|
|
"viewer_url": terminal_projection(self.viewer_url, max_string=1024),
|
|
"error": terminal_projection(self.error, max_string=2 * 1024),
|
|
}
|
|
return bounded_state_projection(state)
|
|
|
|
def collection(self, name: str) -> list[dict[str, Any]]:
|
|
"""Return one bounded terminal projection with stable item identities."""
|
|
if name == "agents":
|
|
return [
|
|
{
|
|
key: terminal_projection(agent.get(key), max_string=256, max_items=5)
|
|
for key in (
|
|
"id",
|
|
"name",
|
|
"parent_id",
|
|
"status",
|
|
"error_message",
|
|
"created_at",
|
|
"updated_at",
|
|
)
|
|
if key in agent
|
|
}
|
|
for agent in self.live_view.agents.values()
|
|
]
|
|
if name == "events":
|
|
return [collection_item_projection(event) for event in self.live_view.events]
|
|
if name == "vulnerabilities":
|
|
reports = (
|
|
self.report_state.vulnerability_reports if self.report_state is not None else []
|
|
)[-MAX_TERMINAL_VULNERABILITIES:]
|
|
result: list[dict[str, Any]] = []
|
|
for index, report in enumerate(reports):
|
|
projected = collection_item_projection(report)
|
|
report_id = projected.get("id")
|
|
if not isinstance(report_id, str) or not report_id:
|
|
projected["id"] = f"vulnerability-{index}"
|
|
result.append(projected)
|
|
return result
|
|
raise ValueError(f"Unknown collection: {name}")
|
|
|
|
def collection_snapshot(self, name: str) -> tuple[int | None, list[dict[str, Any]]]:
|
|
"""Return a collection cursor and complete bounded projection."""
|
|
if name == "events":
|
|
cursor, events = self.live_view.event_snapshot(limit=MAX_TERMINAL_EVENTS)
|
|
return cursor, [collection_item_projection(event) for event in events]
|
|
return None, self.collection(name)
|
|
|
|
def collection_changes(
|
|
self,
|
|
name: str,
|
|
cursor: int,
|
|
) -> tuple[int, list[dict[str, Any]]]:
|
|
"""Return event upserts since a monotonic source cursor."""
|
|
if name != "events":
|
|
raise ValueError(f"Collection {name!r} does not expose incremental changes")
|
|
next_cursor, events = self.live_view.event_changes_since(cursor)
|
|
return next_cursor, [
|
|
collection_item_projection(event) for event in events[-MAX_TERMINAL_EVENTS:]
|
|
]
|
|
|
|
async def handle(self, command: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
handlers = {
|
|
"setup.add_target": self._add_target,
|
|
"setup.set_instruction": self._set_instruction,
|
|
"setup.start": self._start,
|
|
"setup.confirm_mount": self._confirm_mount,
|
|
"agent.send_message": self._send_message,
|
|
"agent.stop": self._stop_agent,
|
|
"viewer.open": self._open_viewer,
|
|
"app.quit": self._quit,
|
|
}
|
|
handler = handlers.get(command)
|
|
if handler is None:
|
|
raise ValueError(f"Unknown command: {command}")
|
|
result = await handler(payload)
|
|
self.notify_changed()
|
|
return result
|
|
|
|
async def _add_target(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
self._require_setup_mutable()
|
|
target = self._required_string(payload, "target")
|
|
if target not in self.targets:
|
|
self.targets.append(target)
|
|
return {"target": target, "total": len(self.targets)}
|
|
|
|
async def _set_instruction(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
self._require_setup_mutable()
|
|
instruction = payload.get("instruction", "")
|
|
if not isinstance(instruction, str):
|
|
raise TypeError("instruction must be a string")
|
|
self.instruction = instruction.strip()
|
|
return {"instruction": self.instruction}
|
|
|
|
async def _start(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
if self.scan_started or self._start_in_progress:
|
|
raise RuntimeError("Scan is already starting or running")
|
|
# A bare prompt launches optimistically, like a coding agent: it skips
|
|
# the network model preflight and surfaces any model error live. A named
|
|
# target keeps the preflight so a real scan does not commit blind.
|
|
verify = payload.get("verify", True)
|
|
if not isinstance(verify, bool):
|
|
raise TypeError("verify must be a boolean")
|
|
# Launching with no target mounts the working directory, so it requires
|
|
# the user's explicit confirmation rather than happening silently.
|
|
mount_working_dir = payload.get("mount_working_dir", False)
|
|
if not isinstance(mount_working_dir, bool):
|
|
raise TypeError("mount_working_dir must be a boolean")
|
|
model = (load_settings().llm.model or "").strip()
|
|
if not model:
|
|
raise ValueError("No model configured. Set STRIX_LLM first.")
|
|
if self._on_start is None:
|
|
raise RuntimeError("Scan start is unavailable")
|
|
if not self.targets:
|
|
if not mount_working_dir:
|
|
raise ValueError("No target set. Add a target first.")
|
|
# Mounting the working directory needs the user's confirmation, and
|
|
# that is asked in the live view. Enter it now and prepare nothing
|
|
# until the answer arrives, so declining leaves no run behind.
|
|
self.pending_workspace_mount = str(Path.cwd())
|
|
self._pending_verify = verify
|
|
self.setup_mode = False
|
|
self.scan_started = True
|
|
self.scan_state = "preparing"
|
|
return {"started": True}
|
|
await self._begin_scan(verify)
|
|
return {"started": True}
|
|
|
|
async def _begin_scan(self, verify: bool) -> None:
|
|
if self._on_start is None:
|
|
raise RuntimeError("Scan start is unavailable")
|
|
self._start_in_progress = True
|
|
try:
|
|
await self._on_start(verify)
|
|
finally:
|
|
self._start_in_progress = False
|
|
self.setup_mode = False
|
|
self.scan_started = True
|
|
self.scan_state = "running"
|
|
|
|
async def _confirm_mount(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Answer the pending working-directory mount asked for in the live view."""
|
|
mount = self.pending_workspace_mount
|
|
if mount is None:
|
|
raise RuntimeError("No mount confirmation is pending")
|
|
approved = payload.get("approved")
|
|
if not isinstance(approved, bool):
|
|
raise TypeError("approved must be a boolean")
|
|
self.pending_workspace_mount = None
|
|
if not approved:
|
|
# Nothing was prepared, so return to the start screen untouched.
|
|
self.workspace_mount = None
|
|
self.enter_setup()
|
|
return {"approved": False}
|
|
self.workspace_mount = mount
|
|
await self._begin_scan(self._pending_verify)
|
|
return {"approved": True}
|
|
|
|
async def _send_message(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
agent_id = self._required_string(payload, "agent_id")
|
|
message = self._required_string(payload, "message")
|
|
if self.coordinator is None:
|
|
raise RuntimeError("Agent coordinator is unavailable")
|
|
if self.scan_loop is None or self.scan_loop.is_closed():
|
|
raise RuntimeError("Scan loop is not ready")
|
|
self.live_view.record_user_message(agent_id, message)
|
|
if self.scan_loop is asyncio.get_running_loop():
|
|
delivered = await self.coordinator.send(
|
|
agent_id,
|
|
{"from": "user", "content": message, "type": "instruction"},
|
|
)
|
|
else:
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
self.coordinator.send(
|
|
agent_id,
|
|
{"from": "user", "content": message, "type": "instruction"},
|
|
),
|
|
self.scan_loop,
|
|
)
|
|
delivered = await asyncio.wrap_future(future)
|
|
if not delivered:
|
|
raise RuntimeError("Message could not be delivered")
|
|
return {"sent": True}
|
|
|
|
async def _stop_agent(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
agent_id = self._required_string(payload, "agent_id")
|
|
agent = self.live_view.agents.get(agent_id)
|
|
if agent is None:
|
|
raise ValueError(f"Unknown agent: {agent_id}")
|
|
status = str(agent.get("status", ""))
|
|
if status not in _STOPPABLE_AGENT_STATUSES:
|
|
raise RuntimeError(f"Agent '{agent_id}' cannot be stopped while {status or 'unknown'}")
|
|
if self.coordinator is None or self.scan_loop is None or self.scan_loop.is_closed():
|
|
raise RuntimeError("Scan loop is not ready")
|
|
if self.scan_loop is asyncio.get_running_loop():
|
|
accepted = await self.coordinator.cancel_descendants_graceful(agent_id)
|
|
else:
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
self.coordinator.cancel_descendants_graceful(agent_id), self.scan_loop
|
|
)
|
|
accepted = await asyncio.wrap_future(future)
|
|
if not accepted:
|
|
raise RuntimeError(f"Agent '{agent_id}' is no longer active")
|
|
return {"stopped": True}
|
|
|
|
async def _open_viewer(self, _payload: dict[str, Any]) -> dict[str, Any]:
|
|
if self.viewer_url:
|
|
with contextlib.suppress(Exception):
|
|
webbrowser.open(self.viewer_url)
|
|
return {"status": "running", "url": self.viewer_url}
|
|
if self.report_state is None:
|
|
self.viewer_status = "failed"
|
|
return {"status": self.viewer_status, "error": "Scan output is not ready"}
|
|
try:
|
|
from strix.interface.tui.backend.messages import (
|
|
send_user_message_to_agent,
|
|
)
|
|
from strix.interface.viewer.server import (
|
|
authorized_url,
|
|
bundle_is_built,
|
|
serve,
|
|
)
|
|
|
|
if not bundle_is_built():
|
|
self.viewer_status = "unavailable"
|
|
return {"status": self.viewer_status, "error": "Viewer UI not built"}
|
|
|
|
def steer(agent_id: str, message: str) -> bool:
|
|
return send_user_message_to_agent(
|
|
coordinator=self.coordinator,
|
|
loop=self.scan_loop,
|
|
live_view=self.live_view,
|
|
target_agent_id=agent_id,
|
|
message=message,
|
|
notify_changed=self.notify_changed,
|
|
wait_for_delivery=True,
|
|
)
|
|
|
|
httpd, url, token = serve(
|
|
self.report_state.get_run_dir(),
|
|
open_browser=True,
|
|
steer_handler=steer,
|
|
)
|
|
self._viewer_httpd = httpd
|
|
self.viewer_url = authorized_url(url, token)
|
|
self.viewer_status = "running"
|
|
with contextlib.suppress(Exception):
|
|
from strix.telemetry import posthog
|
|
|
|
live = self.report_state.run_record.get("status") not in {
|
|
"completed",
|
|
"stopped",
|
|
"failed",
|
|
"interrupted",
|
|
}
|
|
posthog.viewer_opened(source="tui", live=live)
|
|
except Exception: # noqa: BLE001 - viewer startup failures must not crash the TUI
|
|
self.viewer_status = "failed"
|
|
return {"status": self.viewer_status, "error": "Viewer failed to start"}
|
|
else:
|
|
return {"status": self.viewer_status, "url": self.viewer_url}
|
|
|
|
def close_viewer(self) -> None:
|
|
httpd = self._viewer_httpd
|
|
if httpd is None:
|
|
return
|
|
self._viewer_httpd = None
|
|
with contextlib.suppress(Exception):
|
|
httpd.shutdown()
|
|
httpd.server_close()
|
|
|
|
async def _quit(self, _payload: dict[str, Any]) -> dict[str, Any]:
|
|
self.close_viewer()
|
|
if self._on_quit is not None:
|
|
await self._on_quit()
|
|
self.scan_state = "stopped"
|
|
return {"quitting": True}
|
|
|
|
@staticmethod
|
|
def _required_string(payload: dict[str, Any], name: str) -> str:
|
|
value = payload.get(name)
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise ValueError(f"{name} must be a non-empty string")
|
|
return value.strip()
|
|
|
|
def _require_setup_mutable(self) -> None:
|
|
if not self.setup_mode or self.scan_started or self._start_in_progress:
|
|
raise RuntimeError("Setup can no longer be changed after the scan starts")
|