ilia
12a6e2b6bc
FastAPI + Kali tools UI for authorized scans of levkin.ca / LAN targets. Gitignore local data/.env; add pytest + gitleaks CI.
275 lines
10 KiB
Python
275 lines
10 KiB
Python
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)
|