scheduling and jobs, new dataclasses and such better UDP handling

This commit is contained in:
2025-10-17 16:49:30 -05:00
parent 9956667c8f
commit 41306801ae
13 changed files with 771 additions and 169 deletions

View File

@@ -2,11 +2,12 @@ from __future__ import annotations
import os
import subprocess
import xml.etree.ElementTree as ET
import tempfile
from pathlib import Path
from typing import Iterable, List, Dict, Optional, Tuple
from typing import Iterable, List, Dict, Optional, Tuple, Union
from utils.models import HostResult, PortFinding
from utils.scan_config_loader import ScanConfigFile,ScanTarget
class nmap_scanner:
@@ -14,9 +15,11 @@ class nmap_scanner:
UDP_REPORT_PATH = Path() / "data" / "nmap-udp-results.xml"
NMAP_RESULTS_PATH = Path() / "data" / "nmap-results.xml"
def __init__(self, targets:Iterable[str],scan_udp=False):
self.targets = list(targets)
self.scan_udp = scan_udp
def __init__(self, config:ScanConfigFile):
self.scan_config = config
self.targets = config.scan_targets
self.target_list = [t.ip for t in config.scan_targets]
self.scan_udp = config.scan_options.udp_scan
pass
def scan_targets(self):
@@ -24,7 +27,7 @@ class nmap_scanner:
if self.scan_udp:
udp_results = self.run_nmap_udp()
all_results = List[HostResult] = self.merge_host_results(tcp_results,udp_results)
all_results: List[HostResult] = self.merge_host_results(tcp_results,udp_results)
else:
all_results = tcp_results
@@ -35,7 +38,7 @@ class nmap_scanner:
Run a TCP SYN scan across all ports (0-65535) for the given targets and parse results.
Returns a list of HostResult objects.
"""
targets_list = self.targets
targets_list = self.target_list
if not targets_list:
return []
@@ -54,35 +57,135 @@ class nmap_scanner:
self._run_nmap(cmd)
return self.parse_nmap_xml(self.TCP_REPORT_PATH)
def run_nmap_udp(self, ports: Optional[Iterable[int]] = None, min_rate: int = 500, assume_up: bool = True) -> List[HostResult]:
def run_nmap_udp(
self,
ports: Optional[Iterable[int]] = None,
min_rate: int = 500,
assume_up: bool = True,
) -> List[HostResult]:
"""
Run a UDP scan for the provided ports (recommended to keep this list small).
If 'ports' is None, nmap defaults to its "top" UDP ports; full -p- UDP is very slow.
Run UDP scans.
Behavior:
- If `ports` is provided -> single nmap run against all targets using that port list.
- If `ports` is None ->
* For hosts with `expected_udp` defined and non-empty: scan only those ports.
* For hosts with no `expected_udp` (or empty): omit `-p` so nmap uses its default top UDP ports.
Hosts sharing the same explicit UDP port set are grouped into one nmap run.
Returns:
Merged List[HostResult] across all runs.
"""
targets_list = self.targets
targets_list = getattr(self, "target_list", [])
if not targets_list:
return []
cmd = [
"nmap",
"-sU", # UDP scan
"-T3", # less aggressive timing by default for UDP
"--min-rate", str(min_rate),
"-oX", str(self.UDP_REPORT_PATH),
]
if assume_up:
cmd.append("-Pn")
# Optional logger (don't fail if not present)
logger = getattr(self, "logger", None)
def _log(msg: str) -> None:
if logger:
logger.info(msg)
else:
print(msg)
# Case 1: caller provided a global port list -> one run, all targets
if ports:
# Explicit port set
port_list = sorted(set(int(p) for p in ports))
port_list = sorted({int(p) for p in ports})
port_str = ",".join(str(p) for p in port_list)
with tempfile.NamedTemporaryFile(prefix="nmap_udp_", suffix=".xml", delete=False) as tmp:
report_path = tmp.name
cmd = [
"nmap",
"-sU",
"-T3",
"--min-rate", str(min_rate),
"-oX", str(report_path),
]
if assume_up:
cmd.append("-Pn")
cmd.extend(["-p", port_str])
cmd.extend(targets_list)
cmd.extend(targets_list)
_log(f"UDP scan (global ports): {port_str} on {len(targets_list)} host(s)")
self._run_nmap(cmd)
results = self.parse_nmap_xml(report_path)
try:
os.remove(report_path)
except OSError:
pass
return results
self._run_nmap(cmd)
return self.parse_nmap_xml(self.UDP_REPORT_PATH)
# Case 2: per-host behavior using self.scan_config.scan_targets
# Build per-IP port tuple (empty tuple => use nmap's default top UDP ports)
ip_to_ports: Dict[str, Tuple[int, ...]] = {}
# Prefer the IPs present in self.target_list (order/selection comes from there)
# Map from ScanConfigFile / ScanTarget
cfg_targets = getattr(getattr(self, "scan_config", None), "scan_targets", []) or []
# Build quick lookup from config
conf_map: Dict[str, List[int]] = {}
for t in cfg_targets:
# Support either dataclass (attrs) or dict-like
ip = getattr(t, "ip", None) if hasattr(t, "ip") else t.get("ip")
if not ip:
continue
raw_udp = getattr(t, "expected_udp", None) if hasattr(t, "expected_udp") else t.get("expected_udp", [])
conf_map[ip] = list(raw_udp or [])
for ip in targets_list:
raw = conf_map.get(ip, [])
if raw:
ip_to_ports[ip] = tuple(sorted(int(p) for p in raw))
else:
ip_to_ports[ip] = () # empty => use nmap defaults (top UDP ports)
# Group hosts by identical port tuple
groups: Dict[Tuple[int, ...], List[str]] = {}
for ip, port_tuple in ip_to_ports.items():
groups.setdefault(port_tuple, []).append(ip)
all_result_sets: List[List[HostResult]] = []
for port_tuple, ips in groups.items():
# Per-group report path
with tempfile.NamedTemporaryFile(prefix="nmap_udp_", suffix=".xml", delete=False) as tmp:
report_path = tmp.name
cmd = [
"nmap",
"-sU",
"-T3",
"--min-rate", str(min_rate),
"-oX", str(report_path),
]
if assume_up:
cmd.append("-Pn")
if port_tuple:
# explicit per-group ports
port_str = ",".join(str(p) for p in port_tuple)
cmd.extend(["-p", port_str])
_log(f"UDP scan (explicit ports {port_str}) on {len(ips)} host(s): {', '.join(ips)}")
else:
# no -p -> nmap defaults to its top UDP ports
_log(f"UDP scan (nmap top UDP ports) on {len(ips)} host(s): {', '.join(ips)}")
cmd.extend(ips)
self._run_nmap(cmd)
result = self.parse_nmap_xml(report_path)
all_result_sets.append(result)
try:
os.remove(report_path)
except OSError:
pass
if not all_result_sets:
return []
# Merge per-run results into final list
return self.merge_host_results(*all_result_sets)
def merge_host_results(self, *result_sets: List[HostResult]) -> List[HostResult]:
"""
@@ -174,6 +277,6 @@ class nmap_scanner:
self.TCP_REPORT_PATH.unlink()
if self.UDP_REPORT_PATH.exists():
self.UDP_REPORT_PATH.unlink()
if self.NMAP_RESULTS_PATH.exists:
self.NMAP_RESULTS_PATH.unlink()
# if self.NMAP_RESULTS_PATH.exists:
# self.NMAP_RESULTS_PATH.unlink()