analyze performance upgrade

django_project/IoTAnalyzer/analyze/security_analysis.py

@@ -5,15 +5,43 @@ from __future__ import annotations
 import json
 import logging
 import os
+import threading
+import time
+import traceback
 
 from django.conf import settings
-from analyzer import AnalysisEngine, extract_urls, extract_hashes
-from analyzer.extract_pyshark import build_context_from_pcap
+from django.db import connection
+from analyzer import AnalysisEngine, extract_urls, extract_hashes, __version__ as analyzer_version
+from analyzer.extract_tshark import build_context_from_pcap
 from analyzer.findings import Severity
+from analyzer import logutil
 from IoTAnalyzer.models import FindingReport, UrlRaw, HashRaw
 
 logger = logging.getLogger(__name__)
 
+
+def analysis_log_path(filename: str) -> str | None:
+    """Per-capture log written next to the upload."""
+    path = resolve_capture_path(filename)
+    if not path:
+        return None
+    return path + ".analysis.log"
+
+
+def _error_summary(exc: BaseException) -> str:
+    """Short, user-facing description of a failure with the failing location."""
+    tb = traceback.extract_tb(exc.__traceback__)
+    where = f" ({tb[-1].filename.rsplit('/', 1)[-1]}:{tb[-1].lineno})" if tb else ""
+    return f"{type(exc).__name__}: {exc}{where}"[:1000]
+
+
+def _set_error(file_id: int, message: str | None) -> None:
+    from IoTAnalyzer.models import File
+    try:
+        File.objects.filter(id=file_id).update(error_message=message)
+    except Exception as exc:
+        logger.warning("Could not store error message for file=%s: %s", file_id, exc)
+
 _CANDIDATE_DIRS = [
     str(settings.CAPTURE_UPLOAD_DIR),
 ]
@@ -109,38 +137,103 @@ def _persist_artifacts(ctx, raw_dataset_id: int) -> None:
         logger.exception("Hash persistence failed: %s", exc)
 
 
-def analyze_capture(filename: str, file_id: int, raw_dataset_id: int | None = None):
-    """Analyze a capture once and persist findings plus extracted artifacts."""
+def analyze_capture(filename: str, file_id: int, raw_dataset_id: int | None = None,
+                    packet_limit: int | None = None, device_ip: str | None = None):
+    """Analyze a capture once and persist findings plus extracted artifacts.
+
+    ``device_ip`` is the address entered at upload time. It focuses the report on
+    the device under test; when omitted it is looked up from the File record.
+    Raises on failure so the caller can record the reason.
+    """
     path = resolve_capture_path(filename)
     if not path:
-        logger.warning("Capture file not found: %s", filename)
-        return None
+        logger.error("Capture file not found: %s (searched: %s)", filename, _CANDIDATE_DIRS)
+        raise FileNotFoundError(f"Capture file {filename} not found in {_CANDIDATE_DIRS}")
 
-    try:
-        ctx = build_context_from_pcap(path)
-        result = AnalysisEngine().analyze_context(ctx)
-    except Exception as exc:
-        logger.exception("Security analysis failed for %s: %s", filename, exc)
-        return None
+    if packet_limit is None:
+        packet_limit = getattr(settings, "CAPTURE_PACKET_LIMIT", 0)
 
-    try:
-        # Reanalysis replaces the previous report atomically at the file level.
-        FindingReport.objects.filter(file_id=file_id).delete()
-        FindingReport.objects.create(
-            file_id=file_id,
-            result_json=json.dumps(result.to_dict()),
-            finding_count=len(result.findings),
-            max_severity=_max_severity(result),
-        )
-        logger.info("Stored finding report: file_id=%s, findings=%s",
-                    file_id, len(result.findings))
-    except Exception as exc:
-        logger.exception("Finding report persistence failed: %s", exc)
+    if device_ip is None:
+        try:
+            from IoTAnalyzer.models import File
+            device_ip = File.objects.filter(id=file_id).values_list("ip_address", flat=True).first()
+        except Exception:
+            device_ip = None
+
+    started = time.time()
+    logger.info("Capture analysis start: file_id=%s path=%s device_ip=%s packet_limit=%s analyzer=%s",
+                file_id, path, device_ip, packet_limit or "unlimited", analyzer_version)
+    ctx = build_context_from_pcap(path, packet_limit=packet_limit or None, device_ip=device_ip)
+    if ctx.truncated:
+        logger.warning("Capture %s analyzed partially: %s", file_id, ctx.truncation_reason)
+    result = AnalysisEngine().analyze_context(ctx)
+
+    # Reanalysis replaces the previous report atomically at the file level.
+    payload = json.dumps(result.to_dict())
+    FindingReport.objects.filter(file_id=file_id).delete()
+    FindingReport.objects.create(
+        file_id=file_id,
+        result_json=payload,
+        finding_count=len(result.findings),
+        max_severity=_max_severity(result),
+    )
+    logger.info("Stored finding report: file_id=%s findings=%s max_severity=%s json=%d bytes",
+                file_id, len(result.findings), _max_severity(result), len(payload))
 
     _persist_artifacts(ctx, raw_dataset_id)
-
+    logger.info("Capture analysis finished: file_id=%s in %.1fs", file_id, time.time() - started)
     return result
 
 
 def run_security_analysis(filename: str, file_id: int):
-    return analyze_capture(filename, file_id, raw_dataset_id=None)
+    try:
+        return analyze_capture(filename, file_id, raw_dataset_id=None)
+    except Exception as exc:
+        logger.exception("Security analysis failed for %s: %s", filename, exc)
+        return None
+
+
+def start_capture_analysis(file_id: int, filename: str, raw_dataset_id: int | None = None,
+                           synchronous: bool = False) -> None:
+    """Analyze a capture off the request thread and keep File.status current.
+
+    A multi-hundred-megabyte capture takes far longer than any sane HTTP
+    timeout, so the upload response must not wait for the analysis.
+    """
+    if synchronous:
+        _run_capture_analysis(file_id, filename, raw_dataset_id, close_conn=False)
+        return
+    thread = threading.Thread(
+        target=_run_capture_analysis,
+        args=(file_id, filename, raw_dataset_id),
+        daemon=True,
+        name=f"capture-analysis-{file_id}",
+    )
+    thread.start()
+
+
+def _run_capture_analysis(file_id: int, filename: str, raw_dataset_id: int | None = None,
+                          close_conn: bool = True) -> None:
+    from IoTAnalyzer.models import File
+
+    status = "analysis_failed"
+    error = None
+    log_path = analysis_log_path(filename)
+    try:
+        with logutil.run_log(log_path):
+            _set_error(file_id, None)
+            result = analyze_capture(filename, file_id, raw_dataset_id)
+            status = "analyzed" if result is not None else "analysis_failed"
+    except Exception as exc:
+        error = _error_summary(exc)
+        logger.exception("Capture analysis failed for file=%s: %s", file_id, error)
+    finally:
+        try:
+            File.objects.filter(id=file_id).update(status=status, error_message=error)
+            logger.info("Capture %s completed with status %s%s", file_id, status,
+                        f" (log: {log_path})" if log_path else "")
+        except Exception as exc:
+            logger.exception("Could not persist status for file=%s: %s", file_id, exc)
+        # Close only the worker thread's connection, never the request connection.
+        if close_conn:
+            connection.close()