scheduling and jobs, new dataclasses and such better UDP handling

app/main.py

@@ -1,41 +1,34 @@
 #!/usr/bin/env python3
-"""
-port_checker.py
-- expects `expected.json` in same dir (see format below)
-- writes nmap XML to a temp file, parses, compares, prints a report
-"""
-import os
-import json
-import subprocess
-import tempfile
-from datetime import datetime
-import xml.etree.ElementTree as ET
+import logging
+logging.basicConfig(level=logging.INFO)
+
+# TODO:
+# LOGGING
+# TLS SCANNING
+# TLS Version PROBE
+# EMAIL
+
+import time
 from pathlib import Path
 from typing import Dict, List, Set
 
+
+from utils.scan_config_loader import ScanConfigRepository, ScanConfigFile
+from utils.schedule_manager import ScanScheduler
 from utils.scanner import nmap_scanner
 from utils.models import HostResult
-from reporting_jinja import write_html_report_jinja
 
-EXPECTED_FILE = Path() / "data" / "expected.json"
+from reporting_jinja import write_html_report_jinja
+from utils.settings import get_settings
+from utils.common import get_common_utils
+
+logger = logging.getLogger(__file__)
+
+utils = get_common_utils()
+settings = get_settings()
+
 HTML_REPORT_FILE = Path() / "data" / "report.html"
 
-def load_expected(path: Path) -> Dict[str, Dict[str, Set[int]]]:
-    with path.open() as fh:
-        arr = json.load(fh)
-    out = {}
-    for entry in arr:
-        ip = entry["ip"]
-        out[ip] = {
-            "expected_tcp": set(entry.get("expected_tcp", [])),
-            "expected_udp": set(entry.get("expected_udp", [])),
-        }
-    return out
-
-# def write_targets(expected: Dict[str, Dict[str, Set[int]]], path: Path) -> None:
-    path.write_text("\n".join(sorted(expected.keys())) + "\n")
-
-# 
 def results_to_open_sets(
         results: List[HostResult],
         count_as_open: Set[str] = frozenset({"open", "open|filtered"})) -> Dict[str, Dict[str, Set[int]]]:
@@ -56,11 +49,14 @@ def results_to_open_sets(
 
 # Build the "reports" dict (what the HTML renderer expects)
 def build_reports(
-    expected: Dict[str, Dict[str, Set[int]]],
+    scan_config: "ScanConfigFile",
     discovered: Dict[str, Dict[str, Set[int]]],
 ) -> Dict[str, Dict[str, List[int]]]:
     """
-    Create the per-IP delta structure:
+    Create the per-IP delta structure using expected ports from `scan_config.scan_targets`
+    and discovered ports from `discovered`.
+
+    Output format:
     {
       ip: {
         "unexpected_tcp": [...],
@@ -69,15 +65,52 @@ def build_reports(
         "missing_udp": [...]
       }
     }
+
+    Notes:
+      - If a host has no expected UDP ports in the config, `expected_udp` is empty here.
+        (This function reflects *expectations*, not what to scan. Your scan logic can still
+        choose 'top UDP ports' for those hosts.)
+      - The `discovered` dict is expected to use keys "tcp" / "udp" per host.
     """
-    reports: Dict[str, Dict[str, List[int]]] = {}
+    # Build `expected` from scan_config.scan_targets
+    expected: Dict[str, Dict[str, Set[int]]] = {}
+    cfg_targets = getattr(scan_config, "scan_targets", []) or []
+
+    for t in cfg_targets:
+        # Works whether ScanTarget is a dataclass or a dict-like object
+        ip = getattr(t, "ip", None) if hasattr(t, "ip") else t.get("ip")
+        if not ip:
+            continue
+
+        raw_tcp = getattr(t, "expected_tcp", None) if hasattr(t, "expected_tcp") else t.get("expected_tcp", [])
+        raw_udp = getattr(t, "expected_udp", None) if hasattr(t, "expected_udp") else t.get("expected_udp", [])
+
+        exp_tcp = set(int(p) for p in (raw_tcp or []))
+        exp_udp = set(int(p) for p in (raw_udp or []))
+
+        expected[ip] = {
+            "expected_tcp": exp_tcp,
+            "expected_udp": exp_udp,
+        }
+
+    # Union of IPs present in either expectations or discoveries
     all_ips = set(expected.keys()) | set(discovered.keys())
 
+    reports: Dict[str, Dict[str, List[int]]] = {}
     for ip in sorted(all_ips):
+        # Expected sets (default to empty sets if not present)
         exp_tcp = expected.get(ip, {}).get("expected_tcp", set())
         exp_udp = expected.get(ip, {}).get("expected_udp", set())
-        disc_tcp = discovered.get(ip, {}).get("tcp", set())
-        disc_udp = discovered.get(ip, {}).get("udp", set())
+
+        # Discovered sets (default to empty sets if not present)
+        disc_tcp = discovered.get(ip, {}).get("tcp", set()) or set()
+        disc_udp = discovered.get(ip, {}).get("udp", set()) or set()
+
+        # Ensure sets in case caller provided lists
+        if not isinstance(disc_tcp, set):
+            disc_tcp = set(disc_tcp)
+        if not isinstance(disc_udp, set):
+            disc_udp = set(disc_udp)
 
         reports[ip] = {
             "unexpected_tcp": sorted(disc_tcp - exp_tcp),
@@ -85,24 +118,55 @@ def build_reports(
             "unexpected_udp": sorted(disc_udp - exp_udp),
             "missing_udp": sorted(exp_udp - disc_udp),
         }
+
     return reports
 
-def main():
-
-    # repo = ScanConfigRepository()
-
-    if not EXPECTED_FILE.exists():
-        print("Expected File not found")
-        return
-
-    expected = load_expected(EXPECTED_FILE)
-    targets = sorted(expected.keys())
-    scanner = nmap_scanner(targets)
+def run_repo_scan(scan_config:ScanConfigFile):
+    logger.info(f"Starting scan for {scan_config.name}")
+    logger.info("Options: udp=%s tls_sec=%s tls_exp=%s",
+                scan_config.scan_options.udp_scan,
+                scan_config.scan_options.tls_security_scan,
+                scan_config.scan_options.tls_exp_check)
+    logger.info("Targets: %d hosts", len(scan_config.scan_targets))
+    scanner = nmap_scanner(scan_config)
     scan_results = scanner.scan_targets()
     discovered_sets = results_to_open_sets(scan_results, count_as_open={"open", "open|filtered"})
-    reports = build_reports(expected, discovered_sets)
+    reports = build_reports(scan_config, discovered_sets)
     write_html_report_jinja(reports=reports,host_results=scan_results,out_path=HTML_REPORT_FILE,title="Compliance Report",only_issues=True)
     scanner.cleanup()
 
+def main():
+    logger.info(f"{settings.app.name} - v{settings.app.version_major}.{settings.app.version_minor} Started")
+    logger.info(f"Application Running Production flag set to: {settings.app.production}")
+
+    # timezone validation
+    if utils.TextUtils.is_valid_timezone(settings.app.timezone):
+        logger.info(f"Timezone set to {settings.app.timezone}")
+        app_timezone = settings.app.timezone
+    else:
+        logger.warning(f"The Timezone {settings.app.timezone} is invalid, Defaulting to UTC")
+        app_timezone = "America/Danmarkshavn" # UTC
+
+    # load / configure the scan repos
+    repo = ScanConfigRepository()
+    scan_configs = repo.load_all()
+
+    # if in prod - run the scheduler like normal
+    if settings.app.production:
+        sched = ScanScheduler(timezone=app_timezone)
+        sched.start()
+
+        jobs = sched.schedule_configs(scan_configs, run_scan_fn=run_repo_scan)
+        logger.info("Scheduled %d job(s).", jobs)
+
+        try:
+            while True:
+                time.sleep(3600)
+        except KeyboardInterrupt:
+            sched.shutdown()
+    else:
+        # run single scan in dev mode
+        run_repo_scan(scan_configs[0])
+    
 if __name__ == "__main__":
     main()