ilia
12a6e2b6bc
FastAPI + Kali tools UI for authorized scans of levkin.ca / LAN targets. Gitignore local data/.env; add pytest + gitleaks CI.
96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
import asyncio
|
|
import time
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
import structlog
|
|
|
|
from app.config import settings
|
|
from app.models import ToolResult
|
|
|
|
log = structlog.get_logger(__name__)
|
|
|
|
# Track running processes for cancellation: scan_id -> list[asyncio.subprocess.Process]
|
|
_running: dict[str, list[asyncio.subprocess.Process]] = {}
|
|
|
|
|
|
def register_process(scan_id: str, proc: asyncio.subprocess.Process) -> None:
|
|
_running.setdefault(scan_id, []).append(proc)
|
|
|
|
|
|
def unregister_process(scan_id: str, proc: asyncio.subprocess.Process) -> None:
|
|
procs = _running.get(scan_id, [])
|
|
if proc in procs:
|
|
procs.remove(proc)
|
|
|
|
|
|
async def cancel_scan_processes(scan_id: str) -> None:
|
|
for proc in _running.get(scan_id, []):
|
|
try:
|
|
proc.kill()
|
|
except ProcessLookupError:
|
|
pass
|
|
_running[scan_id] = []
|
|
|
|
|
|
async def run_tool(
|
|
command: list[str],
|
|
scan_id: str | None = None,
|
|
cwd: str | None = None,
|
|
timeout: int | None = None,
|
|
on_line: Callable[[str], Awaitable[None]] | None = None,
|
|
) -> ToolResult:
|
|
"""Spawn process without shell; stream stdout line-by-line."""
|
|
timeout = timeout or settings.tool_timeout_seconds
|
|
start = time.monotonic()
|
|
stdout_parts: list[str] = []
|
|
stderr_parts: list[str] = []
|
|
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*command,
|
|
cwd=cwd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
if scan_id:
|
|
register_process(scan_id, proc)
|
|
|
|
async def read_stream(stream: asyncio.StreamReader, is_stderr: bool) -> None:
|
|
while True:
|
|
line = await stream.readline()
|
|
if not line:
|
|
break
|
|
text = line.decode(errors="replace").rstrip("\n")
|
|
if is_stderr:
|
|
stderr_parts.append(text)
|
|
else:
|
|
stdout_parts.append(text)
|
|
if on_line:
|
|
await on_line(text)
|
|
|
|
try:
|
|
await asyncio.wait_for(
|
|
asyncio.gather(
|
|
read_stream(proc.stdout, False), # type: ignore[arg-type]
|
|
read_stream(proc.stderr, True), # type: ignore[arg-type]
|
|
),
|
|
timeout=timeout,
|
|
)
|
|
exit_code = await proc.wait()
|
|
except asyncio.TimeoutError:
|
|
proc.kill()
|
|
await proc.wait()
|
|
exit_code = -1
|
|
stderr_parts.append(f"timeout after {timeout}s")
|
|
log.warning("tool_timeout", command=command, timeout=timeout)
|
|
finally:
|
|
if scan_id:
|
|
unregister_process(scan_id, proc)
|
|
|
|
duration = time.monotonic() - start
|
|
return ToolResult(
|
|
stdout="\n".join(stdout_parts),
|
|
stderr="\n".join(stderr_parts),
|
|
exit_code=exit_code,
|
|
duration=duration,
|
|
)
|