diff --git a/.env.example b/.env.example index f1b4686..b87eea6 100644 --- a/.env.example +++ b/.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 diff --git a/README.md b/README.md index e12b742..0adf7a9 100644 --- a/README.md +++ b/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//` 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 | |----------|--------|-------------| diff --git a/pyproject.toml b/pyproject.toml index c45cfbb..da64d65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "rich>=13.0", "pyyaml>=6.0", "jinja2>=3.0", + "python-dotenv>=1.0.0", ] [project.optional-dependencies] diff --git a/src/deepzero/cli.py b/src/deepzero/cli.py index 355f56e..e767081 100644 --- a/src/deepzero/cli.py +++ b/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) diff --git a/src/deepzero/engine/runner.py b/src/deepzero/engine/runner.py index fd7853b..0c7e332 100644 --- a/src/deepzero/engine/runner.py +++ b/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()