poc2_worked

app/rules.py

@@ -8,7 +8,10 @@ import shutil
 import signal
 import subprocess
 import tempfile
+import tarfile
 import threading
+import uuid
+from copy import deepcopy
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
@@ -28,7 +31,7 @@ class RuleManager:
     MAX_THRESHOLD_BYTES = 256 * 1024
     SOURCE_INDEX_URL = "https://www.openinfosecfoundation.org/rules/index.yaml"
     DEFAULT_SOURCE = "et/open"
-    SOURCE_NAME_RE = re.compile(r"^[A-Za-z0-9_.+-]+/[A-Za-z0-9_.+-]+$")
+    SOURCE_NAME_RE = re.compile(r"^[A-Za-z0-9_.+-]+(?:/[A-Za-z0-9_.+-]+)?$")
 
     def __init__(
         self,
@@ -42,7 +45,23 @@ class RuleManager:
         self._lock = threading.RLock()
         self._operation_lock = threading.RLock()
         self._update_lock = threading.Lock()
+        self._source_queue_lock = threading.RLock()
+        self._source_queue = {
+            "id": "",
+            "status": "idle",
+            "phase": "idle",
+            "created_at": None,
+            "started_at": None,
+            "finished_at": None,
+            "total": 0,
+            "completed": 0,
+            "failed": 0,
+            "message": "No queued source operation",
+            "items": [],
+        }
         self._last_result = "not changed"
+        self._snapshot_dir = Path(self.config.suricata_custom_rules).parent / "rule-snapshots"
+        self._snapshot_dir.mkdir(parents=True, exist_ok=True)
         self._ensure_files()
 
     def _ensure_files(self) -> None:
@@ -56,10 +75,12 @@ class RuleManager:
         threshold = self._read(self.config.suricata_threshold_config)
         with self._lock:
             last_result = self._last_result
-        vendor_rules = "/var/lib/suricata/rules/suricata.rules"
+        vendor_root = self.config.suricata_persist_lib_dir
+        vendor_rules = os.path.join(vendor_root, "rules", "suricata.rules")
         source_index = _first_existing_path(
-            "/var/lib/suricata/update/cache/index.yaml",
-            "/var/lib/suricata/rules/cache/index.yaml",
+            os.path.join(vendor_root, "rules", ".cache", "index.yaml"),
+            os.path.join(vendor_root, "update", "cache", "index.yaml"),
+            os.path.join(vendor_root, "rules", "cache", "index.yaml"),
         )
         return {
             "available": self.suricata_available,
@@ -77,6 +98,7 @@ class RuleManager:
             "source_index_updated_at": _file_mtime_iso(source_index) if source_index else None,
             "source_index_url": self.SOURCE_INDEX_URL,
             "last_result": last_result,
+            "snapshots": len(self.list_snapshots()),
         }
 
     def content(self) -> dict:
@@ -139,6 +161,134 @@ class RuleManager:
             current += line + "\n"
             return self.replace_threshold_config(current)
 
+    def add_threshold(
+        self,
+        sid: int,
+        *,
+        threshold_type: str = "limit",
+        track: str = "by_src",
+        count: int = 5,
+        seconds: int = 60,
+    ) -> RuleActionResult:
+        sid = int(sid)
+        threshold_type = str(threshold_type or "limit").strip().lower()
+        track = str(track or "by_src").strip().lower()
+        count = max(1, min(100000, int(count)))
+        seconds = max(1, min(86400, int(seconds)))
+        if sid <= 0:
+            return RuleActionResult(False, "SID must be a positive integer")
+        if threshold_type not in {"limit", "threshold", "both"}:
+            return RuleActionResult(False, "threshold type must be limit, threshold or both")
+        if track not in {"by_src", "by_dst", "by_rule", "by_both", "by_flow"}:
+            return RuleActionResult(False, "unsupported threshold tracker")
+        line = f"threshold gen_id 1, sig_id {sid}, type {threshold_type}, track {track}, count {count}, seconds {seconds}"
+        with self._operation_lock:
+            current = self._read(self.config.suricata_threshold_config)
+            if line.casefold() in {x.strip().casefold() for x in current.splitlines() if x.strip()}:
+                return RuleActionResult(True, f"SID {sid} already has that threshold")
+            if current and not current.endswith("\n"):
+                current += "\n"
+            current += line + "\n"
+            return self.replace_threshold_config(current)
+
+    def create_snapshot(self, reason: str = "manual") -> RuleActionResult:
+        try:
+            with self._operation_lock:
+                path = self._create_snapshot(reason)
+            return RuleActionResult(True, f"rule snapshot created: {path.name}")
+        except Exception as exc:
+            return RuleActionResult(False, f"could not create rule snapshot: {exc}")
+
+    def list_snapshots(self) -> list[dict]:
+        out = []
+        try:
+            paths = sorted(self._snapshot_dir.glob("rules-*.tar.gz"), key=lambda p: p.stat().st_mtime, reverse=True)
+        except OSError:
+            return []
+        for path in paths[:20]:
+            try:
+                stat = path.stat()
+            except OSError:
+                continue
+            out.append({
+                "id": path.name,
+                "created_at": datetime.fromtimestamp(stat.st_mtime, timezone.utc).isoformat(),
+                "size_bytes": int(stat.st_size),
+            })
+        return out
+
+    def rollback_snapshot(self, snapshot_id: str) -> RuleActionResult:
+        name = os.path.basename(str(snapshot_id or ""))
+        if not re.fullmatch(r"rules-[A-Za-z0-9_.-]+\.tar\.gz", name):
+            return RuleActionResult(False, "invalid rule snapshot")
+        path = self._snapshot_dir / name
+        if not path.is_file():
+            return RuleActionResult(False, "rule snapshot not found")
+        if not self.suricata_available:
+            return RuleActionResult(False, "Suricata is not available in this mode")
+        with self._operation_lock:
+            backup = self._create_snapshot("pre-rollback")
+            try:
+                with tempfile.TemporaryDirectory(prefix="rules-rollback-") as td:
+                    root = Path(td)
+                    with tarfile.open(path, "r:gz") as tar:
+                        for member in tar.getmembers():
+                            dest = (root / member.name).resolve()
+                            if root.resolve() not in dest.parents and dest != root.resolve():
+                                raise ValueError("unsafe snapshot path")
+                        tar.extractall(root)
+                    custom = (root / "custom.rules").read_text(encoding="utf-8") if (root / "custom.rules").exists() else ""
+                    threshold = (root / "threshold.config").read_text(encoding="utf-8") if (root / "threshold.config").exists() else ""
+                    validation = self.validate(custom, threshold)
+                    if not validation.ok:
+                        return RuleActionResult(False, f"snapshot validation failed: {validation.message}")
+                    self._atomic_write(self.config.suricata_custom_rules, custom)
+                    self._atomic_write(self.config.suricata_threshold_config, threshold)
+                    vendor = root / "vendor.rules"
+                    if vendor.exists():
+                        vendor_dest = Path(self._suricata_update_data_dir()) / "rules" / "suricata.rules"
+                        vendor_dest.parent.mkdir(parents=True, exist_ok=True)
+                        shutil.copyfile(vendor, vendor_dest)
+                    sources = root / "sources"
+                    if sources.exists():
+                        source_dest = Path(self._suricata_update_data_dir()) / "update" / "sources"
+                        if source_dest.exists():
+                            shutil.rmtree(source_dest)
+                        shutil.copytree(sources, source_dest)
+                result = self.reload()
+                if result.ok:
+                    return RuleActionResult(True, f"restored {name}; {result.message}; safety snapshot {backup.name}")
+                return RuleActionResult(False, f"restored files but {result.message}; safety snapshot {backup.name}")
+            except Exception as exc:
+                return RuleActionResult(False, f"rollback failed: {exc}; safety snapshot {backup.name}")
+
+    def _create_snapshot(self, reason: str) -> Path:
+        self._snapshot_dir.mkdir(parents=True, exist_ok=True)
+        stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+        safe_reason = re.sub(r"[^A-Za-z0-9_.-]+", "-", str(reason or "snapshot"))[:40].strip("-") or "snapshot"
+        target = self._snapshot_dir / f"rules-{stamp}-{safe_reason}-{uuid.uuid4().hex[:6]}.tar.gz"
+        with tarfile.open(target, "w:gz") as tar:
+            for source, arcname in (
+                (Path(self.config.suricata_custom_rules), "custom.rules"),
+                (Path(self.config.suricata_threshold_config), "threshold.config"),
+                (Path(self._suricata_update_data_dir()) / "rules" / "suricata.rules", "vendor.rules"),
+            ):
+                if source.is_file():
+                    tar.add(source, arcname=arcname, recursive=False)
+            sources = Path(self._suricata_update_data_dir()) / "update" / "sources"
+            if sources.is_dir():
+                tar.add(sources, arcname="sources", recursive=True)
+        self._prune_snapshots(12)
+        return target
+
+    def _prune_snapshots(self, keep: int) -> None:
+        paths = sorted(self._snapshot_dir.glob("rules-*.tar.gz"), key=lambda p: p.stat().st_mtime, reverse=True)
+        for path in paths[max(1, int(keep)):]:
+            try:
+                path.unlink()
+            except OSError:
+                pass
+
     def update_vendor_rules(self) -> RuleActionResult:
         if not self.suricata_available:
             return RuleActionResult(False, "Suricata rule updates are unavailable in this mode")
@@ -167,9 +317,13 @@ class RuleManager:
         enabled_proc = self._run_suricata_update(["list-sources", "--enabled"], timeout=30)
         enabled = _parse_enabled_sources(enabled_proc.stdout or "") if enabled_proc.returncode == 0 else set()
         sources = _parse_source_catalog(catalog.stdout or "")
+        default_replaced = any(
+            source.get("name") in enabled and self.DEFAULT_SOURCE in source.get("replaces", [])
+            for source in sources
+        )
         for source in sources:
             source["default"] = source["name"] == self.DEFAULT_SOURCE
-            source["enabled"] = source["default"] or source["name"] in enabled
+            source["enabled"] = source["name"] in enabled or (source["default"] and not default_replaced)
             source["can_toggle"] = not source["default"] and not bool(source.get("parameters"))
         return {
             "ok": True,
@@ -180,6 +334,8 @@ class RuleManager:
             "enabled_sources": sorted(
                 {source["name"] for source in sources if source.get("enabled")}
             ),
+            "data_dir": self._suricata_update_data_dir(),
+            "queue": self.source_queue_status(),
             "status": self.status(),
         }
 
@@ -247,7 +403,155 @@ class RuleManager:
         finally:
             self._update_lock.release()
 
+    def queue_sources(self, source_names: list[str]) -> RuleActionResult:
+        if not self.suricata_available:
+            return RuleActionResult(False, "Suricata rule sources are unavailable in this mode")
+
+        normalized: list[str] = []
+        seen: set[str] = set()
+        for raw in source_names or []:
+            name = str(raw or "").strip()
+            if not name or name in seen:
+                continue
+            if not self.SOURCE_NAME_RE.fullmatch(name):
+                return RuleActionResult(False, f"invalid rule source name: {name}")
+            seen.add(name)
+            normalized.append(name)
+        if not normalized:
+            return RuleActionResult(False, "select at least one rule source")
+        if len(normalized) > 128:
+            return RuleActionResult(False, "too many rule sources in one queue (maximum 128)")
+
+        with self._source_queue_lock:
+            if self._source_queue.get("status") in {"queued", "running"}:
+                return RuleActionResult(False, "a rule-source download queue is already running")
+            job_id = uuid.uuid4().hex[:12]
+            now = datetime.now(timezone.utc).isoformat()
+            self._source_queue = {
+                "id": job_id,
+                "status": "queued",
+                "phase": "waiting",
+                "created_at": now,
+                "started_at": None,
+                "finished_at": None,
+                "total": len(normalized),
+                "completed": 0,
+                "failed": 0,
+                "message": f"Queued {len(normalized)} source(s)",
+                "items": [
+                    {"source": name, "status": "pending", "message": "Waiting"}
+                    for name in normalized
+                ],
+            }
+
+        worker = threading.Thread(
+            target=self._source_queue_worker,
+            args=(job_id, normalized),
+            name=f"rule-source-queue-{job_id}",
+            daemon=True,
+        )
+        worker.start()
+        return RuleActionResult(True, f"Queued {len(normalized)} rule source(s) for sequential download")
+
+    def source_queue_status(self) -> dict:
+        with self._source_queue_lock:
+            return deepcopy(self._source_queue)
+
+    def _source_queue_worker(self, job_id: str, source_names: list[str]) -> None:
+        self._queue_job_update(job_id, status="running", phase="catalog", started_at=datetime.now(timezone.utc).isoformat(), message="Loading persistent source catalog")
+        self._update_lock.acquire()
+        try:
+            catalog = self.source_catalog()
+            if not catalog.get("ok"):
+                self._queue_job_finish(job_id, "failed", str(catalog.get("error") or "could not read source catalog"))
+                return
+            by_name = {str(item.get("name")): item for item in catalog.get("sources", [])}
+            changed = 0
+            failed = 0
+            completed = 0
+
+            for index, name in enumerate(source_names):
+                self._queue_item_update(job_id, index, "running", "Enabling source")
+                source = by_name.get(name)
+                if source is None:
+                    failed += 1
+                    self._queue_item_update(job_id, index, "failed", "Source is not present in the free OISF catalog")
+                    self._queue_job_update(job_id, failed=failed)
+                    continue
+                if source.get("parameters"):
+                    failed += 1
+                    params = ", ".join(source.get("parameters") or [])
+                    self._queue_item_update(job_id, index, "failed", f"Requires parameters: {params}")
+                    self._queue_job_update(job_id, failed=failed)
+                    continue
+                if source.get("enabled"):
+                    completed += 1
+                    self._queue_item_update(job_id, index, "done", "Already enabled; will refresh with active feeds")
+                    self._queue_job_update(job_id, completed=completed)
+                    continue
+
+                proc = self._run_suricata_update(["enable-source", name], timeout=90)
+                if proc.returncode != 0:
+                    failed += 1
+                    self._queue_item_update(job_id, index, "failed", _command_tail(proc.stdout, "enable-source failed"))
+                    self._queue_job_update(job_id, failed=failed)
+                    continue
+                changed += 1
+                completed += 1
+                self._queue_item_update(job_id, index, "done", "Enabled in persistent /data source state")
+                self._queue_job_update(job_id, completed=completed)
+
+            self._queue_job_update(
+                job_id,
+                phase="download",
+                message=f"Downloading and merging all active feeds ({completed} selected source(s) ready)",
+            )
+            update_result = self._run_vendor_update_unlocked()
+            if not update_result.ok:
+                self._queue_job_finish(job_id, "failed", update_result.message)
+                return
+
+            final_status = "partial" if failed else "completed"
+            summary = f"{completed} source(s) ready, {failed} failed; {update_result.message}"
+            if changed == 0 and failed == 0:
+                summary = f"Selected sources were already enabled; {update_result.message}"
+            self._queue_job_finish(job_id, final_status, summary)
+        except Exception as exc:
+            self._queue_job_finish(job_id, "failed", f"rule-source queue failed: {exc}")
+        finally:
+            self._update_lock.release()
+
+    def _queue_job_update(self, job_id: str, **fields) -> None:
+        with self._source_queue_lock:
+            if self._source_queue.get("id") != job_id:
+                return
+            self._source_queue.update(fields)
+
+    def _queue_item_update(self, job_id: str, index: int, status: str, message: str) -> None:
+        with self._source_queue_lock:
+            if self._source_queue.get("id") != job_id:
+                return
+            items = self._source_queue.get("items") or []
+            if 0 <= index < len(items):
+                items[index]["status"] = status
+                items[index]["message"] = str(message)[:1000]
+
+    def _queue_job_finish(self, job_id: str, status: str, message: str) -> None:
+        self._queue_job_update(
+            job_id,
+            status=status,
+            phase="done",
+            finished_at=datetime.now(timezone.utc).isoformat(),
+            message=str(message)[:1600],
+        )
+        with self._lock:
+            self._last_result = str(message)[:1600]
+
     def _run_vendor_update_unlocked(self) -> RuleActionResult:
+        try:
+            self._create_snapshot("pre-update")
+        except Exception as exc:
+            print(f"[rules] snapshot before update failed: {exc}", flush=True)
         try:
             proc = subprocess.run(
                 ["/opt/ids/scripts/update-rules.sh"],
@@ -269,11 +573,14 @@ class RuleManager:
             self._last_result = result.message
         return result
 
-    @staticmethod
-    def _run_suricata_update(args: list[str], timeout: int) -> subprocess.CompletedProcess:
+    def _suricata_update_data_dir(self) -> str:
+        return str(getattr(self.config, "suricata_persist_lib_dir", "/data/lib/suricata"))
+
+    def _run_suricata_update(self, args: list[str], timeout: int) -> subprocess.CompletedProcess:
+        command = ["suricata-update", *args, "-D", self._suricata_update_data_dir()]
         try:
             return subprocess.run(
-                ["suricata-update", *args],
+                command,
                 check=False,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT,
@@ -282,7 +589,7 @@ class RuleManager:
             )
         except (OSError, subprocess.TimeoutExpired) as exc:
             return subprocess.CompletedProcess(
-                ["suricata-update", *args],
+                command,
                 127,
                 stdout=f"suricata-update could not run: {exc}",
             )
@@ -321,6 +628,11 @@ class RuleManager:
                 name = os.path.basename(source)
                 shutil.copyfile(source, os.path.join(rules_dir, name))
                 copied.add(name)
+            state_dir = os.path.dirname(self.config.suricata_custom_rules)
+            for source in glob.glob(os.path.join(state_dir, "*.lst")):
+                if not os.path.isfile(source):
+                    continue
+                shutil.copyfile(source, os.path.join(rules_dir, os.path.basename(source)))
             local_name = os.path.basename(self.config.suricata_local_rules) or "local.rules"
             if local_name not in copied and os.path.isfile(self.config.suricata_local_rules):
                 shutil.copyfile(self.config.suricata_local_rules, os.path.join(rules_dir, local_name))
@@ -331,6 +643,8 @@ class RuleManager:
                 "-T",
                 "-c",
                 self.config.suricata_config,
+                "--include",
+                self.config.suricata_output_config,
                 "-l",
                 log_dir,
                 "-s",
@@ -339,6 +653,12 @@ class RuleManager:
                 f"vars.address-groups.HOME_NET={self.config.suricata_home_net}",
                 "--set",
                 f"threshold-file={threshold_path}",
+                "--set",
+                "app-layer.protocols.tls.ja3-fingerprints=yes",
+                "--set",
+                "app-layer.protocols.tls.ja4-fingerprints=yes",
+                "--set",
+                "app-layer.protocols.ssh.hassh=yes",
             ]
             try:
                 proc = subprocess.run(
@@ -380,6 +700,10 @@ class RuleManager:
                     self._last_result = validation.message
                 return validation
 
+            try:
+                self._create_snapshot(f"pre-{label.replace(' ', '-')}")
+            except Exception as exc:
+                print(f"[rules] snapshot before {label} change failed: {exc}", flush=True)
             self._atomic_write(path, content)
             reload_result = self.reload()
             if reload_result.ok:
@@ -467,7 +791,7 @@ def _parse_source_catalog(output: str) -> list[dict]:
 def _parse_enabled_sources(output: str) -> set[str]:
     result: set[str] = set()
     for raw in _strip_ansi(output).splitlines():
-        match = re.match(r"^\s*-\s+([A-Za-z0-9_.+-]+/[A-Za-z0-9_.+-]+)\s*$", raw)
+        match = re.match(r"^\s*-\s+([A-Za-z0-9_.+-]+(?:/[A-Za-z0-9_.+-]+)?)\s*$", raw)
         if match:
             result.add(match.group(1))
     return result