Refactor/cli telemetry and metrics

.env.example

@@ -12,6 +12,8 @@
 # OPENAI_API_KEY=sk-...
 # ANTHROPIC_API_KEY=sk-ant-...
 # OPENROUTER_API_KEY=sk-or-v1-...
+# VERTEXAI_PROJECT=
+# VERTEXAI_LOCATION=
 
 # =========================================================
 # 2. GHIDRA Headless Configuration

README.md

@@ -32,7 +32,6 @@
 
 - [⚡️ Quickstart](#️-quickstart)
 - [⚙️ Architecture](#️-architecture)
-- [📦 Installation](#-installation)
 - [🚀 CLI Reference](#-cli)
 - [📝 Pipeline Syntax](#-pipeline-syntax)
 - [🛠 Shipped Processors](#-shipped-processors)
@@ -46,20 +45,33 @@
 
 DeepZero requires a target corpus of files to analyze and a pipeline configuration detailing how to process them. We provide a complete example pipeline designed to hunt for vulnerabilities in Windows kernel drivers using the LOLDrivers dataset as a blocklist.
 
-1. **Download a sample corpus**
+1. **Clone & Install**
+   DeepZero requires **Python 3.11+**.
+   ```bash
+   git clone https://github.com/416rehman/DeepZero.git
+   cd DeepZero
+   pip install -e .
+   ```
+
+2. **Configure Environment**
+   If you intend to use the AI analysis stages, configure your API keys by creating a `.env` file from the provided layout:
+   ```bash
+   cp .env.example .env
+   ```
+
+3. **Download a sample corpus**
    For this example, we will use the open-source Snappy Driver Installer (SDI) driver pack which contains thousands of real-world Windows drivers:
    👉 **[https://sdi-tool.org/download/](https://sdi-tool.org/download/)**
 
    Download the standard/torrent package and extract the archive to a local path (e.g., `C:\drivers`).
 
-2. **Run the LOLDrivers pipeline**
-   Pass your driver path to the CLI, alongside the provided pipeline definition:
+4. **Run the LOLDrivers pipeline**
+   Pass your driver path to the CLI, alongside the provided pipeline configuration file:
    ```bash
    deepzero run C:\drivers -p .\pipelines\loldrivers\pipeline.yaml
    ```
 
-3. **Monitor progress and state**
-   DeepZero will safely parallelize execution, cache intermediate outputs, and show a live terminal dashboard. If you need to stop unexpectedly, simply press `Ctrl+C`. Re-run the exact same command later, and DeepZero will instantly resume from disk state.
+   *DeepZero will safely parallelize execution, cache intermediate outputs, and show a live terminal dashboard. If you need to stop unexpectedly, simply press `Ctrl+C`. Re-run the exact same command later, and DeepZero will instantly resume from disk state.*
 
 ---
 
@@ -128,44 +140,6 @@ Sees all active samples at once. Returns a list of sample IDs to keep, in the de
 
 ---
 
-## 📦 Installation
-
-Requires **Python 3.11+**.
-
-```bash
-git clone https://github.com/416rehman/DeepZero.git
-cd DeepZero
-
-# install with all optional extras
-pip install -e ".[full]"
-
-# or install only what you need
-pip install -e ".[llm]"    # LLM support via litellm
-pip install -e ".[pe]"     # PE header parsing via lief
-pip install -e ".[serve]"  # REST API server (starlette + uvicorn)
-
-# copy and populate environment variables
-cp .env.example .env
-```
-
-### Core Dependencies
-
-| Package | Purpose |
-|---------|---------|
-| `click` | CLI framework |
-| `rich` | Terminal UI, progress bars, live dashboard |
-| `pyyaml` | Pipeline YAML parsing |
-| `jinja2` | LLM prompt templating |
-
-### Optional Dependencies
-
-| Extra | Package | Purpose |
-|-------|---------|---------|
-| `llm` | `litellm` | LLM provider abstraction (OpenAI, Anthropic, Vertex AI, etc.) |
-| `pe` | `lief` | PE header parsing for the `pe_ingest` processor |
-| `serve` | `starlette`, `uvicorn` | REST API server |
-
----
 
 ## 🚀 CLI
 
@@ -211,9 +185,12 @@ deepzero serve --host 127.0.0.1 --port 8420 -w work/
 
 Resume is automatic. If a `work/<pipeline>/` directory exists with prior state, `deepzero run` resumes from where it left off. Use `--clean` to discard and restart.
 
-### `deepzero serve`
+### `deepzero serve` (⚠️ WIP / Experimental)
+
+Starts a read-only REST API (Starlette + Uvicorn) for querying run state and sample data. 
 
-Starts a read-only REST API (Starlette + Uvicorn) for querying run state and sample data.
+> [!WARNING]
+> The REST API server is currently a Work-In-Progress (WIP) and is partially broken/unstable.
 
 | Endpoint | Method | Description |
 |----------|--------|-------------|

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "rich>=13.0",
     "pyyaml>=6.0",
     "jinja2>=3.0",
+    "python-dotenv>=1.0.0",
 ]
 
 [project.optional-dependencies]

src/deepzero/cli.py

@@ -9,6 +9,7 @@
 import click
 from rich.console import Console
 from rich.logging import RichHandler
+from rich.markup import escape
 from rich.table import Table
 
 from deepzero.engine.types import RunStatus
@@ -24,6 +25,17 @@
     }
 )
 
+_LOG_COLORS: tuple[str, ...] = (
+    "cyan",
+    "magenta",
+    "green",
+    "yellow",
+    "bright_cyan",
+    "bright_magenta",
+    "bright_green",
+    "bright_yellow",
+)
+
 
 class _ShortNameFormatter(logging.Formatter):
     # strips deepzero prefix to keep log lines short and scannable
@@ -44,8 +56,16 @@ def format(self, record: logging.LogRecord) -> str:
             record.exc_info = None
             record.exc_text = None
 
-        record.msg = f"{short:>20} | {record.msg}"
-        return super().format(record)
+        msg = super().format(record)
+
+        msg_escaped = escape(msg)
+
+        import zlib
+
+        color = _LOG_COLORS[zlib.crc32(short.encode("utf-8")) % len(_LOG_COLORS)]
+
+        # format dynamically and override the final payload that RichHandler receives
+        return f"[{color}]\\[{short}][/{color}] {msg_escaped}"
 
 
 def _setup_logging(verbose: bool) -> None:
@@ -56,7 +76,10 @@ def _setup_logging(verbose: bool) -> None:
         show_path=False,
         show_level=False,
         show_time=True,
+        log_time_format="%H:%M:%S",
+        markup=True,
     )
+    # the formatter now handles the full string construction and markup tagging
     handler.setFormatter(_ShortNameFormatter("%(message)s"))
     logging.basicConfig(level=level, handlers=[handler])
     if not verbose:
@@ -208,14 +231,20 @@ def status(pipeline: str | None, work_dir: str | None, verbose: bool):
     """show current pipeline run status"""
     from deepzero.engine.state import StateStore
 
+    _setup_logging(verbose)
     if work_dir:
         work_path = Path(work_dir)
     elif pipeline:
         import deepzero.stages  # noqa: F401
         from deepzero.engine.pipeline import load_pipeline
 
-        _setup_logging(False)
-        pipeline_def = load_pipeline(pipeline)
+        _load_env()
+        try:
+            pipeline_def = load_pipeline(pipeline)
+        except ValueError as e:
+            console.print(f"[bold red]X ERROR[/]: {e}")
+            raise SystemExit(1)
+
         work_path = pipeline_def.work_dir
     else:
         console.print("[red]specify --pipeline or --work-dir[/]")
@@ -242,10 +271,10 @@ def status(pipeline: str | None, work_dir: str | None, verbose: bool):
     if run_state.completed_at:
         console.print(f"  completed: {run_state.completed_at}")
 
-    _print_stats(run_state)
+    manifest = state_store.load_manifest()
+    _print_stats(run_state, manifest)
 
     # show manifest summary
-    manifest = state_store.load_manifest()
     if manifest:
         verdicts: dict[str, int] = {}
         for entry in manifest:
@@ -262,6 +291,7 @@ def status(pipeline: str | None, work_dir: str | None, verbose: bool):
 def validate(pipeline_ref: str):
     """validate a pipeline definition"""
     _setup_logging(False)
+    _load_env()
 
     import deepzero.stages  # noqa: F401
     from deepzero.engine.pipeline import validate_pipeline
@@ -405,14 +435,7 @@ def interactive(model: str, work_dir: str, verbose: bool):
             response = llm.complete(history)
             history.append({"role": "assistant", "content": response})
             console.print(f"\n[bold cyan]deepzero>[/] {response}\n")
-        except (
-            RuntimeError,
-            ValueError,
-            TypeError,
-            OSError,
-            ConnectionError,
-            TimeoutError,
-        ) as e:
+        except Exception as e:
             console.print(f"[red]llm error ({type(e).__name__}): {e}[/]")
 
 
@@ -429,6 +452,8 @@ def serve(host: str, port: int, work_dir: str):
     console.print(f"[bold cyan]deepzero serve[/] - http://{host}:{port}")
     console.print(f"  work_dir: {work_dir}")
 
+    _load_env()
+
     from deepzero.api.server import create_app
 
     try:
@@ -441,9 +466,14 @@ def serve(host: str, port: int, work_dir: str):
     uvicorn.run(app, host=host, port=port)
 
 
-def _print_stats(run_state) -> None:
-    per_stage = run_state.stats.get("per_stage", {})
-    if not per_stage:
+def _print_stats(run_state, manifest: list[dict[str, Any]] | None = None) -> None:
+    # `manifest` is accepted for caller compatibility, but current manifest
+    # entries do not persist per-stage history, so stage statistics must come
+    # from the run state's cached per-stage counters.
+    _ = manifest
+    per_stage: dict[str, dict[str, int]] = dict(run_state.stats.get("per_stage", {}))
+
+    if not run_state.stages:
         discovered = run_state.stats.get("discovered", 0)
         if discovered:
             console.print(f"  discovered: {discovered}")
@@ -455,13 +485,36 @@ def _print_stats(run_state) -> None:
     table.add_column("filtered", style="yellow", justify="right")
     table.add_column("failed", style="red", justify="right")
 
-    for stage_name, counts in per_stage.items():
-        table.add_row(
-            stage_name,
-            str(counts.get("completed", 0)),
-            str(counts.get("filtered", 0)),
-            str(counts.get("failed", 0)),
-        )
+    for i, stage_name in enumerate(run_state.stages):
+        if i == 0:
+            # first stage is always ingest, stats are held in discovered
+            discovered = run_state.stats.get("discovered", 0)
+            table.add_row(
+                stage_name,
+                str(discovered) if discovered else "[dim white]·[/]",
+                "[dim white]·[/]",
+                "[dim white]·[/]",
+            )
+        else:
+            counts = per_stage.get(stage_name, {})
+            if not counts:
+                # stage completely unstarted (force dim white to strip column colors)
+                table.add_row(
+                    f"[dim white]◦ {stage_name}[/]",
+                    "[dim white]·[/]",
+                    "[dim white]·[/]",
+                    "[dim white]·[/]",
+                )
+            else:
+                p = counts.get("completed", 0)
+                f = counts.get("filtered", 0)
+                err = counts.get("failed", 0)
+                table.add_row(
+                    stage_name,
+                    str(p) if p else "[dim white]·[/]",
+                    str(f) if f else "[dim white]·[/]",
+                    str(err) if err else "[dim white]·[/]",
+                )
 
     console.print(table)
 

src/deepzero/engine/runner.py

@@ -104,6 +104,7 @@ def run(
         target: Path,
         run_state: RunState,
     ) -> RunState:
+        self._active_run_state = run_state
         self._install_signal_handler()
         run_state.mark_running()
         self.state_store.save_run(run_state)
@@ -112,7 +113,7 @@ def run(
             return self._execute_pipeline_stages(target, run_state)
         except KeyboardInterrupt:
             log.warning("interrupted by user - saving state")
-            run_state.status = RunStatus.INTERRUPTED
+            run_state.mark_interrupted()
             self.state_store.save_run(run_state)
             return run_state
         except PROCESSOR_ERRORS as e:
@@ -159,6 +160,7 @@ def _run_all_stages(
         stage_names: list[str],
     ) -> RunState:
         sample_states = self._resume_or_ingest(target, run_state, stage_names)
+        self._active_sample_states = sample_states
         if sample_states is None:
             return run_state
 
@@ -836,6 +838,13 @@ def _restore_signal_handler(self) -> None:
     def _handle_signal(self, signum, frame) -> None:
         if self._shutdown_event.is_set():
             log.warning("forced shutdown")
+            if hasattr(self, "_active_run_state") and getattr(self, "_active_run_state", None):
+                self._active_run_state.mark_interrupted()
+                self.state_store.save_run(self._active_run_state)
+            if hasattr(self, "_active_sample_states") and getattr(
+                self, "_active_sample_states", None
+            ):
+                self.state_store.save_manifest(list(self._active_sample_states.values()))
             os._exit(1)
         log.warning("shutdown requested (press ctrl+c again to force)")
         self._shutdown_event.set()