import asyncio import re import uuid from collections import defaultdict from datetime import datetime, timezone from pathlib import Path from typing import Any import httpx import structlog from app import db, llm, profiles, tools from app.config import settings from app.models import Finding, Scan log = structlog.get_logger(__name__) _active_scans: set[str] = set() _target_active: dict[str, set[str]] = defaultdict(set) _cancelled: set[str] = set() _scan_semaphore: asyncio.Semaphore | None = None def _semaphore() -> asyncio.Semaphore: global _scan_semaphore if _scan_semaphore is None: _scan_semaphore = asyncio.Semaphore(settings.max_concurrent_scans) return _scan_semaphore class WsBroadcaster: def __init__(self) -> None: self._queues: list[asyncio.Queue[dict[str, Any]]] = [] def subscribe(self) -> asyncio.Queue[dict[str, Any]]: q: asyncio.Queue[dict[str, Any]] = asyncio.Queue() self._queues.append(q) return q async def send(self, message: dict[str, Any]) -> None: for q in list(self._queues): await q.put(message) _scan_broadcasters: dict[str, WsBroadcaster] = {} def get_broadcaster(scan_id: str) -> WsBroadcaster: if scan_id not in _scan_broadcasters: _scan_broadcasters[scan_id] = WsBroadcaster() return _scan_broadcasters[scan_id] def is_scan_active(scan_id: str) -> bool: return scan_id in _active_scans async def target_busy(target: str) -> tuple[bool, str]: """One queued or running scan per target at a time.""" if await db.has_active_scan_for_target(target): return True, "A scan is already queued or running for this target" return False, "" def _audit_log(target: str, profile: str, source_ip: str) -> None: path = Path(settings.audit_log_path) path.parent.mkdir(parents=True, exist_ok=True) ts = datetime.now(timezone.utc).isoformat() line = f"{ts}\ttarget={target}\tprofile={profile}\tsource_ip={source_ip}\n" with path.open("a", encoding="utf-8") as f: f.write(line) def _extract_open_ports(nmap_output: str) -> str: ports: list[str] = [] for m in re.finditer(r'portid="(\d+)"', nmap_output): ports.append(m.group(1)) if not ports: for m in re.finditer(r"(\d+)/tcp\s+open", nmap_output): ports.append(m.group(1)) return ",".join(sorted(set(ports), key=int)) if ports else "80,443" async def _resolve_final_url(start_url: str) -> str: timeout = httpx.Timeout(connect=5.0, read=10.0, write=5.0, pool=5.0) async with httpx.AsyncClient(follow_redirects=True, timeout=timeout, verify=True) as client: try: r = await client.head(start_url) except Exception: r = await client.get(start_url) return str(r.url) def _ensure_https(url: str) -> str: if url.startswith("https://"): return url if url.startswith("http://"): return "https://" + url.removeprefix("http://") return "https://" + url def _classify_nonzero_phase(pdef: profiles.PhaseDef, output: str) -> str | None: if pdef.name == "dir_enum": low = output.lower() if "no wordlist found" in low or ("wordlist" in low and ("no such file" in low or "not found" in low)): return "inconclusive" if "429" in low or "too many requests" in low: return "blocked" if "403" in low or "forbidden" in low: return "blocked" if "timeout" in low: return "inconclusive" if pdef.name == "ssl_check": low = output.lower() if "timeout" in low: return "inconclusive" return None async def queue_scan(scan: Scan, source_ip: str = "unknown") -> None: """Enqueue scan; waits on semaphore when max concurrent scans are active.""" _audit_log(scan.target, scan.profile, source_ip) asyncio.create_task(_run_with_slot(scan.id)) async def _run_with_slot(scan_id: str) -> None: async with _semaphore(): await _run_scan(scan_id) async def cancel_scan(scan_id: str) -> bool: _cancelled.add(scan_id) await tools.cancel_scan_processes(scan_id) await db.update_scan_status(scan_id, "cancelled") bc = get_broadcaster(scan_id) await bc.send({"type": "done", "scan_id": scan_id, "status": "cancelled"}) scan = await db.get_scan(scan_id) if scan and scan.bulk_id: await db.refresh_bulk_run_status(scan.bulk_id) return True async def _run_scan(scan_id: str) -> None: scan = await db.get_scan(scan_id) if not scan: return _active_scans.add(scan_id) _target_active[scan.target].add(scan_id) bc = get_broadcaster(scan_id) llm_down = False try: await db.update_scan_status(scan_id, "running") if scan.bulk_id: await db.refresh_bulk_run_status(scan.bulk_id) phase_defs = profiles.get_phase_defs(scan.profile) open_ports = "80,443" ctx_base = profiles._ctx(scan.target, scan.options, scan_id) ctx_dynamic: dict[str, str] = {} if scan.options.include_web or scan.options.include_ssl: try: start = profiles._target_url(scan.target) final_url = await _resolve_final_url(start) final_https_url = _ensure_https(final_url) except Exception as exc: final_url = profiles._target_url(scan.target) final_https_url = profiles._https_target_url(scan.target) ctx_dynamic["url_resolution_error"] = str(exc) ctx_dynamic["final_url"] = final_url ctx_dynamic["final_https_url"] = final_https_url phase_rows: list[tuple[str, str, str, str]] = [] for pdef in phase_defs: if not pdef.should_run(scan.options): continue cmd = profiles.render_command( pdef.command_template, {**ctx_base, **ctx_dynamic, "open_ports": open_ports}, ) phase_rows.append((str(uuid.uuid4()), pdef.name, pdef.tool, " ".join(cmd))) phases = await db.create_phases(scan_id, phase_rows) runnable_defs = [pd for pd in phase_defs if pd.should_run(scan.options)] for p in phases: await bc.send({"type": "phase", "phase_id": p.id, "phase": p.name, "status": p.status}) for phase, pdef in zip(phases, runnable_defs): if scan_id in _cancelled: break ctx = {**ctx_base, **ctx_dynamic, "open_ports": open_ports} command = profiles.render_command(pdef.command_template, ctx) if pdef.name == "ssh_brute" and "22" not in (open_ports or "").split(","): note = f"skipped: ssh port 22 not in open ports ({open_ports})" await db.update_phase(phase.id, "skipped", raw_output=note, exit_code=0) await bc.send({"type": "phase", "phase_id": phase.id, "phase": phase.name, "status": "skipped"}) continue await db.update_phase(phase.id, "running", started=True) await bc.send({"type": "phase", "phase_id": phase.id, "phase": phase.name, "status": "running"}) async def on_line(line: str, pid: str = phase.id) -> None: await db.append_phase_output(pid, line) await bc.send({"type": "log", "phase_id": pid, "line": line}) last_result = None attempts = 1 + max(0, pdef.retries) for attempt in range(attempts): if attempt > 0 and pdef.retry_delay_seconds > 0: await asyncio.sleep(pdef.retry_delay_seconds) last_result = await tools.run_tool( command, scan_id=scan_id, on_line=on_line, timeout=pdef.timeout_seconds, ) if last_result.exit_code in pdef.success_exit_codes: break assert last_result is not None result = last_result status = "complete" if result.exit_code in pdef.success_exit_codes else "failed" full_out = result.stdout + ("\n" + result.stderr if result.stderr else "") if status == "failed": classified = _classify_nonzero_phase(pdef, full_out) if classified: status = classified await db.update_phase(phase.id, status, raw_output=full_out, exit_code=result.exit_code) await bc.send( { "type": "phase", "phase_id": phase.id, "phase": phase.name, "status": status, "duration": result.duration, } ) if phase.name == "port_scan" and result.stdout: open_ports = _extract_open_ports(result.stdout) if pdef.llm_analyze and full_out.strip(): phase.raw_output = full_out findings, unavailable = await llm.analyze_phase( phase, full_out, target=scan.target, profile=scan.profile, model=scan.options.llm_model, ) if unavailable: llm_down = True for f in findings: await db.save_finding(f) await bc.send({"type": "finding", "finding": f.model_dump(mode="json")}) final = "cancelled" if scan_id in _cancelled else "complete" await db.update_scan_status(scan_id, final, llm_unavailable=llm_down) await bc.send({"type": "done", "scan_id": scan_id, "status": final, "llm_unavailable": llm_down}) except Exception as exc: log.exception("scan_failed", scan_id=scan_id, error=str(exc)) await db.update_scan_status(scan_id, "failed") await bc.send({"type": "error", "message": str(exc)}) await bc.send({"type": "done", "scan_id": scan_id, "status": "failed"}) finally: _active_scans.discard(scan_id) _target_active[scan.target].discard(scan_id) _cancelled.discard(scan_id) if scan and scan.bulk_id: await db.refresh_bulk_run_status(scan.bulk_id)