✨feature: Enhance scanning capabilities with Docker insights and concurrent processing
This commit is contained in:
43
app/main.py
43
app/main.py
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
@@ -7,9 +8,9 @@ from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import FileResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from .config import DEFAULT_SUBNET
|
||||
from .config import DEFAULT_SUBNET, SCAN_WORKERS
|
||||
from .db import init_db
|
||||
from .scanner import discover_hosts, scan_host
|
||||
from .scanner import HostResult, discover_docker_ports, discover_hosts, merge_docker_ports, scan_host
|
||||
from .service import (
|
||||
complete_scan,
|
||||
create_scan,
|
||||
@@ -77,23 +78,31 @@ def run_scan(subnet: str | None = None) -> dict[str, Any]:
|
||||
try:
|
||||
discovered = discover_hosts(subnet)
|
||||
scan_state.set_total_hosts(len(discovered))
|
||||
docker_ports_by_ip = discover_docker_ports({h["ip"] for h in discovered})
|
||||
|
||||
for idx, host in enumerate(discovered, start=1):
|
||||
scan_state.set_current_host(host["ip"])
|
||||
detailed = scan_host(host["ip"])
|
||||
if not detailed:
|
||||
scan_state.update_progress(idx, host_count)
|
||||
continue
|
||||
host_count += 1
|
||||
if not detailed.hostname and host.get("hostname"):
|
||||
detailed.hostname = host["hostname"]
|
||||
if not detailed.mac and host.get("mac"):
|
||||
detailed.mac = host["mac"]
|
||||
if not detailed.vendor and host.get("vendor"):
|
||||
detailed.vendor = host["vendor"]
|
||||
processed_count = 0
|
||||
max_workers = min(max(1, SCAN_WORKERS), max(1, len(discovered)))
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as pool:
|
||||
future_map = {pool.submit(scan_host, host["ip"], host): host for host in discovered}
|
||||
|
||||
upsert_host(scan_id, detailed)
|
||||
scan_state.update_progress(idx, host_count)
|
||||
for future in as_completed(future_map):
|
||||
seed = future_map[future]
|
||||
scan_state.set_current_host(seed["ip"])
|
||||
processed_count += 1
|
||||
try:
|
||||
detailed = future.result()
|
||||
except Exception:
|
||||
detailed = HostResult(
|
||||
ip=seed["ip"],
|
||||
hostname=seed.get("hostname"),
|
||||
mac=seed.get("mac"),
|
||||
vendor=seed.get("vendor"),
|
||||
)
|
||||
|
||||
detailed = merge_docker_ports(detailed, docker_ports_by_ip.get(detailed.ip, []))
|
||||
upsert_host(scan_id, detailed)
|
||||
host_count += 1
|
||||
scan_state.update_progress(processed_count, host_count)
|
||||
|
||||
mark_missing_devices(scan_id)
|
||||
complete_scan(scan_id, "completed", host_count)
|
||||
|
||||
Reference in New Issue
Block a user