mirror of
https://github.com/416rehman/DeepZero-Agentic-Vulnerability-Research-Pipeline
synced 2026-08-09 11:55:01 +00:00
Merge PR #16: browsable run reports, resume and auth fixes
feat: browse a run's results in the browser, plus resume and auth fixes
This commit is contained in:
@@ -51,3 +51,6 @@ docs/_site/
|
||||
.jekyll-cache
|
||||
docs/.jekyll-cache/
|
||||
_site/
|
||||
|
||||
# local research artifacts: run logs, generated reports, scratch previews
|
||||
research/
|
||||
|
||||
@@ -13,7 +13,21 @@ name: loldrivers
|
||||
description: windows kernel driver vulnerability research pipeline
|
||||
version: "2.0"
|
||||
|
||||
model: claude-code/sonnet
|
||||
model: claude-code/opus
|
||||
|
||||
# how `deepzero report` presents this pipeline's results. every field is
|
||||
# optional - a pipeline that declares nothing still gets a useful report.
|
||||
report:
|
||||
title: Windows kernel driver review
|
||||
entity: driver
|
||||
classification_key: classification
|
||||
vulnerable_when: [vulnerable]
|
||||
safe_when: [safe]
|
||||
columns:
|
||||
- semgrep_scanner.finding_count
|
||||
- discover.priority_score
|
||||
- decompile.function_count
|
||||
- discover.dangerous_imports
|
||||
|
||||
settings:
|
||||
work_dir: work
|
||||
|
||||
@@ -174,8 +174,10 @@ def run(
|
||||
# ensure built-in stages are registered
|
||||
import deepzero.stages # noqa: F401
|
||||
from deepzero.engine.pipeline import load_pipeline
|
||||
from deepzero.engine.report import write_report
|
||||
from deepzero.engine.state import RunState, StateStore
|
||||
|
||||
log_report = logging.getLogger("deepzero.report")
|
||||
target_path = Path(target).resolve()
|
||||
|
||||
try:
|
||||
@@ -246,6 +248,12 @@ def run(
|
||||
if is_resume:
|
||||
run_state = existing_run
|
||||
run_state.status = RunStatus.RUNNING
|
||||
# the pipeline's model may have changed since the run was created; record
|
||||
# what this resumed run will actually call so status/reports don't lie
|
||||
if run_state.model != pipeline_def.model:
|
||||
log_msg = f"model changed since last run: {run_state.model or '(unset)'} -> {pipeline_def.model}"
|
||||
console.print(f"[yellow]![/] {log_msg}")
|
||||
run_state.model = pipeline_def.model
|
||||
else:
|
||||
# initialize fresh state
|
||||
state_store.save_pipeline_snapshot(pipeline_def.raw_yaml)
|
||||
@@ -257,7 +265,79 @@ def run(
|
||||
model=pipeline_def.model,
|
||||
)
|
||||
|
||||
# a live report the user can open straight away and watch fill in. it is
|
||||
# rebuilt when results actually land rather than on a timer, and rate
|
||||
# limited so a fast stage cannot spend the run's time writing html
|
||||
report_dir = pipeline_def.work_dir / "report"
|
||||
report_index = report_dir / "index.html"
|
||||
last_written = [0.0]
|
||||
|
||||
def _refresh_report(force: bool = False) -> None:
|
||||
if not force and time.monotonic() - last_written[0] < 15:
|
||||
return
|
||||
try:
|
||||
write_report(pipeline_def.work_dir, report_dir, config=pipeline_def.report)
|
||||
last_written[0] = time.monotonic()
|
||||
except (OSError, ValueError, KeyError, TypeError) as exc:
|
||||
log_report.debug("could not refresh the report: %s", exc)
|
||||
|
||||
_refresh_report(force=True)
|
||||
console.print(f" report [bold]{report_index.resolve().as_uri()}[/]")
|
||||
console.print(" [dim]open it now - it updates itself as results land[/]\n")
|
||||
|
||||
runner.progress_hook = _refresh_report
|
||||
try:
|
||||
run_state = runner.run(target_path, run_state)
|
||||
finally:
|
||||
_refresh_report(force=True) # settle the final state, no reload banner
|
||||
console.print(f"\n report [bold]{report_index.resolve().as_uri()}[/]")
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option("--pipeline", "-p", default=None, help="pipeline name or path")
|
||||
@click.option("--work-dir", "-w", default=None, help="work directory (overrides --pipeline)")
|
||||
@click.option("--out", "-o", default=None, help="output directory (default <work_dir>/report)")
|
||||
@click.option("--open", "open_browser", is_flag=True, help="open the report when finished")
|
||||
@click.option("--verbose", "-v", is_flag=True, help="verbose logging")
|
||||
def report(
|
||||
pipeline: str | None, work_dir: str | None, out: str | None, open_browser: bool, verbose: bool
|
||||
):
|
||||
"""build a browsable HTML report from a run's results (safe mid-run)"""
|
||||
_setup_logging(verbose)
|
||||
|
||||
from deepzero.engine.report import write_report
|
||||
|
||||
report_cfg: dict = {}
|
||||
if work_dir:
|
||||
work_path = Path(work_dir)
|
||||
elif pipeline:
|
||||
import deepzero.stages # noqa: F401
|
||||
from deepzero.engine.pipeline import load_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
|
||||
report_cfg = pipeline_def.report
|
||||
else:
|
||||
console.print("[bold red]X ERROR[/]: pass --pipeline or --work-dir")
|
||||
raise SystemExit(1)
|
||||
|
||||
if not work_path.exists():
|
||||
console.print(f"[bold red]X ERROR[/]: work directory does not exist: {work_path}")
|
||||
raise SystemExit(1)
|
||||
|
||||
html_path, json_path = write_report(work_path, Path(out) if out else None, config=report_cfg)
|
||||
console.print(f"[green]\\[ok][/] report written to [bold]{html_path}[/]")
|
||||
console.print(f" machine-readable copy: {json_path}")
|
||||
|
||||
if open_browser:
|
||||
import webbrowser
|
||||
|
||||
webbrowser.open(html_path.resolve().as_uri())
|
||||
|
||||
|
||||
@main.command()
|
||||
|
||||
@@ -15,10 +15,16 @@ import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from deepzero.engine.backends.base import BackendError, CLIAgentBackend
|
||||
from deepzero.engine.backends.base import BackendAuthError, BackendError, CLIAgentBackend
|
||||
|
||||
log = logging.getLogger("deepzero.llm.claude_code")
|
||||
|
||||
# what to tell the user when the cli reports an auth failure. deepzero does not
|
||||
# inspect claude's credential store: whatever auth the cli has is the auth we
|
||||
# use, and if it has none the cli says so and we relay it.
|
||||
_AUTH_HINT = "run `claude` in a terminal and sign in, then retry"
|
||||
|
||||
|
||||
# the prompt carries untrusted content (decompiled malware, attacker-controlled
|
||||
# strings). a completion never needs tools, so deny the ones that could turn a
|
||||
# prompt injection into code execution or exfiltration.
|
||||
@@ -43,11 +49,8 @@ class ClaudeCodeBackend(CLIAgentBackend):
|
||||
binary_names: ClassVar[tuple[str, ...]] = ("claude",)
|
||||
install_hint: ClassVar[str] = "install it and sign in (https://claude.com/claude-code)."
|
||||
# never use --bare: it disables oauth/keychain reads, which is exactly the
|
||||
# auth this backend exists to use
|
||||
subscription_auth_env_blocklist: ClassVar[tuple[str, ...]] = (
|
||||
"ANTHROPIC_API_KEY",
|
||||
"ANTHROPIC_AUTH_TOKEN",
|
||||
)
|
||||
# auth this backend exists to use. the cli decides how to authenticate from
|
||||
# its own config and environment - deepzero does not second-guess it.
|
||||
|
||||
def __init__(self, model: str, **kwargs: Any):
|
||||
super().__init__(model, **kwargs)
|
||||
@@ -71,6 +74,13 @@ class ClaudeCodeBackend(CLIAgentBackend):
|
||||
paths.append(home / ".local" / "bin" / "claude.exe")
|
||||
return paths
|
||||
|
||||
def classify_error(self, detail: str, status: Any = None) -> BackendError:
|
||||
err = super().classify_error(detail, status)
|
||||
# the cli's own auth message is the source of truth; just add the remedy
|
||||
if isinstance(err, BackendAuthError):
|
||||
return BackendAuthError(f"{err} - {_AUTH_HINT}")
|
||||
return err
|
||||
|
||||
def build_argv(self, system: str) -> tuple[list[str], str]:
|
||||
argv = [str(self._binary), "-p", "--output-format", "json"]
|
||||
|
||||
|
||||
@@ -38,9 +38,9 @@ class LLMProvider:
|
||||
"""send messages to the llm and return the response text.
|
||||
handles rate limiting with adaptive backoff."""
|
||||
backend = self.backend
|
||||
# forward all options; each backend uses what applies (litellm passes
|
||||
# generation kwargs to the api, cli backends read controls like timeout
|
||||
# and ignore the rest) so e.g. timeout= is not silently dropped
|
||||
# every option is forwarded and each backend takes what applies to it:
|
||||
# litellm passes generation kwargs to the api, while cli backends read
|
||||
# controls such as timeout and ignore the rest
|
||||
merged = {**self.default_kwargs, **kwargs}
|
||||
backoff = initial_backoff
|
||||
|
||||
|
||||
@@ -35,6 +35,7 @@ class PipelineDefinition:
|
||||
stage_specs: list[StageSpec],
|
||||
pipeline_dir: Path,
|
||||
raw_yaml: str,
|
||||
report: dict[str, Any] | None = None,
|
||||
):
|
||||
self.name = name
|
||||
self.description = description
|
||||
@@ -44,6 +45,9 @@ class PipelineDefinition:
|
||||
self.stage_specs = stage_specs
|
||||
self.pipeline_dir = pipeline_dir
|
||||
self.raw_yaml = raw_yaml
|
||||
# optional presentation hints for `deepzero report` - pipelines that
|
||||
# declare nothing still get a useful report
|
||||
self.report = report or {}
|
||||
|
||||
# resolved processor instances
|
||||
self.ingest_processor: IngestProcessor | None = None
|
||||
@@ -138,6 +142,7 @@ def load_pipeline(
|
||||
name=name,
|
||||
description=description,
|
||||
model=model,
|
||||
report=data.get("report") or {},
|
||||
settings=settings,
|
||||
knowledge=knowledge,
|
||||
stage_specs=stage_specs,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -9,7 +9,7 @@ import threading
|
||||
import time
|
||||
import traceback as tb_module
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
@@ -76,6 +76,7 @@ class PipelineRunner:
|
||||
default_max_workers: int = 4,
|
||||
console: Console | None = None,
|
||||
dashboard: PipelineDashboard | None = None,
|
||||
progress_hook: Callable[[], None] | None = None,
|
||||
):
|
||||
self.ingest = ingest
|
||||
self.stages = stages
|
||||
@@ -86,9 +87,20 @@ class PipelineRunner:
|
||||
self.default_max_workers = default_max_workers
|
||||
self.console = console or Console()
|
||||
self.dashboard = dashboard
|
||||
# called when new results have landed, so a caller can refresh a live
|
||||
# view. never allowed to interrupt the run
|
||||
self.progress_hook = progress_hook
|
||||
self._shutdown_event = threading.Event()
|
||||
self._original_sigint = None
|
||||
|
||||
def _notify_progress(self) -> None:
|
||||
if self.progress_hook is None:
|
||||
return
|
||||
try:
|
||||
self.progress_hook()
|
||||
except Exception as exc: # noqa: BLE001 - a view refresh must never fail a run
|
||||
log.debug("progress hook failed: %s", exc)
|
||||
|
||||
def _make_entry(self, state: SampleState) -> ProcessorEntry:
|
||||
# centralizes ProcessorEntry construction for map/reduce/batch
|
||||
return ProcessorEntry(
|
||||
@@ -223,6 +235,8 @@ class PipelineRunner:
|
||||
else:
|
||||
self._run_map(processor, active, spec, stage_stats)
|
||||
|
||||
self._notify_progress()
|
||||
|
||||
self._apply_stage_limit(spec, sample_states, stage_stats)
|
||||
|
||||
elapsed = time.monotonic() - t0
|
||||
@@ -460,6 +474,7 @@ class PipelineRunner:
|
||||
for s in dirty:
|
||||
self.state_store.save_sample(s)
|
||||
dirty.clear()
|
||||
self._notify_progress()
|
||||
for s in dirty:
|
||||
self.state_store.save_sample(s)
|
||||
else:
|
||||
@@ -507,6 +522,7 @@ class PipelineRunner:
|
||||
for s in dirty:
|
||||
self.state_store.save_sample(s)
|
||||
dirty.clear()
|
||||
self._notify_progress()
|
||||
for s in dirty:
|
||||
self.state_store.save_sample(s)
|
||||
|
||||
@@ -522,7 +538,9 @@ class PipelineRunner:
|
||||
|
||||
if result.status == StageStatus.COMPLETED:
|
||||
if result.data and "__skipped" in result.data:
|
||||
state.mark_stage_skipped(spec.name, result.data["__skipped"])
|
||||
# work that was already done counts as passed and the sample
|
||||
# continues to the next stage, per MapProcessor.should_skip
|
||||
state.mark_stage_cached(spec.name, result.data["__skipped"])
|
||||
else:
|
||||
state.mark_stage_completed(
|
||||
spec.name,
|
||||
|
||||
@@ -51,6 +51,9 @@ class StageOutput:
|
||||
# namespaced processor output - never merged across stages
|
||||
data: dict[str, Any] = field(default_factory=dict)
|
||||
error: str | None = None
|
||||
# why this stage did no work (cached output, filtered out). kept separate from
|
||||
# error so a normal skip is never reported as a failure
|
||||
skip_reason: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -105,6 +108,22 @@ class SampleState:
|
||||
self.error = error
|
||||
self.verdict = SampleStatus.FAILED
|
||||
|
||||
def mark_stage_cached(self, stage_name: str, reason: str = "") -> None:
|
||||
"""work for this stage was already done (e.g. output already on disk).
|
||||
|
||||
the sample stays in the pipeline: per MapProcessor.should_skip, a skipped
|
||||
sample counts as passed. this is NOT mark_stage_skipped, which filters the
|
||||
sample out of the run entirely.
|
||||
"""
|
||||
if stage_name not in self.history:
|
||||
self.history[stage_name] = StageOutput()
|
||||
stage = self.history[stage_name]
|
||||
stage.status = StageStatus.COMPLETED
|
||||
stage.verdict = Verdict.CONTINUE
|
||||
stage.completed_at = _now()
|
||||
if reason:
|
||||
stage.skip_reason = reason
|
||||
|
||||
def mark_stage_skipped(self, stage_name: str, reason: str = "") -> None:
|
||||
if stage_name not in self.history:
|
||||
self.history[stage_name] = StageOutput()
|
||||
@@ -113,7 +132,7 @@ class SampleState:
|
||||
stage.verdict = Verdict.FILTER
|
||||
stage.completed_at = _now()
|
||||
if reason:
|
||||
stage.error = reason
|
||||
stage.skip_reason = reason
|
||||
self.verdict = SampleStatus.FILTERED
|
||||
|
||||
def is_stage_done(self, stage_name: str) -> bool:
|
||||
|
||||
@@ -132,14 +132,11 @@ class TestCommandConstruction:
|
||||
|
||||
|
||||
class TestEnvHandling:
|
||||
def test_api_key_hidden_by_default(self):
|
||||
def test_env_is_passed_through_untouched(self):
|
||||
# deepzero does not strip or inject auth vars - the cli decides how to
|
||||
# authenticate from its own config and environment
|
||||
with patch.dict("os.environ", {"ANTHROPIC_API_KEY": "sk-ant-x"}, clear=False):
|
||||
env = _backend().build_env()
|
||||
assert "ANTHROPIC_API_KEY" not in env
|
||||
|
||||
def test_api_key_preserved_when_opted_out(self):
|
||||
with patch.dict("os.environ", {"ANTHROPIC_API_KEY": "sk-ant-x"}, clear=False):
|
||||
env = _backend(prefer_subscription_auth=False).build_env()
|
||||
assert env["ANTHROPIC_API_KEY"] == "sk-ant-x"
|
||||
|
||||
|
||||
@@ -167,7 +164,10 @@ class TestResponseParsing:
|
||||
with patch("subprocess.run", return_value=_completed(body)):
|
||||
with pytest.raises(BackendAuthError) as exc:
|
||||
_backend().raw_complete([{"role": "user", "content": "hi"}])
|
||||
assert "not authenticated" in str(exc.value)
|
||||
msg = str(exc.value)
|
||||
assert "not authenticated" in msg
|
||||
# relay the cli's own message plus a fixed remedy - no credential probing
|
||||
assert "sign in" in msg
|
||||
|
||||
def test_rate_limit_is_classified(self):
|
||||
body = _result_json(is_error=True, api_error_status=429, result="rate limit exceeded")
|
||||
@@ -303,9 +303,12 @@ class TestValidation:
|
||||
errors = self._validate("claude-code")
|
||||
assert errors and "claude code cli not found" in errors[0].lower()
|
||||
|
||||
def test_claude_code_binding_does_not_require_api_keys(self):
|
||||
# no ANTHROPIC_API_KEY in env, yet validation passes
|
||||
with patch.dict("os.environ", {}, clear=True):
|
||||
def test_claude_code_binding_does_not_require_api_keys(self, monkeypatch):
|
||||
# no api keys anywhere, yet a valid subscription credential validates.
|
||||
# (clears only the key vars - wiping the whole env would also remove the
|
||||
# credential location, which is a different failure)
|
||||
monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
|
||||
with patch(_FIND, return_value="/usr/bin/claude"):
|
||||
assert self._validate("claude-code") == []
|
||||
|
||||
|
||||
@@ -0,0 +1,341 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from deepzero.engine.report import (
|
||||
BUCKET_SUSPICIOUS,
|
||||
BUCKET_VULNERABLE,
|
||||
ReportConfig,
|
||||
collect,
|
||||
render_index,
|
||||
write_report,
|
||||
)
|
||||
from deepzero.engine.state import RunState, SampleState, StateStore
|
||||
from deepzero.engine.types import RunStatus
|
||||
|
||||
|
||||
def _seed(tmp_path, *, with_findings=True, with_assessment=True, classification="vulnerable"):
|
||||
"""build a work dir shaped like a real run: one risky driver, one clean."""
|
||||
work = tmp_path / "work" / "loldrivers"
|
||||
store = StateStore(work)
|
||||
store.save_run(
|
||||
RunState(
|
||||
run_id="run_1",
|
||||
pipeline="loldrivers",
|
||||
target="C:/drivers",
|
||||
model="claude-code/opus",
|
||||
status=RunStatus.COMPLETED,
|
||||
)
|
||||
)
|
||||
|
||||
risky = SampleState(
|
||||
sample_id="aaa1",
|
||||
sha256="a" * 64,
|
||||
filename="risky.sys",
|
||||
source_path="C:/drivers/risky.sys",
|
||||
)
|
||||
risky.mark_stage_completed(
|
||||
"discover", data={"priority_score": 8.0, "dangerous_imports": ["MmMapIoSpace"]}
|
||||
)
|
||||
risky.mark_stage_completed(
|
||||
"decompile",
|
||||
artifacts={"ghidra_result": "decompiled/ghidra_result.json"},
|
||||
data={"device_name": "RiskyDev", "function_count": 42, "ioctl_count": 2},
|
||||
)
|
||||
if with_findings:
|
||||
risky.mark_stage_completed("semgrep_scanner", data={"finding_count": 2})
|
||||
if with_assessment:
|
||||
risky.mark_stage_completed(
|
||||
"assess",
|
||||
artifacts={"llm_output": "assessment.md"},
|
||||
data={"classification": classification},
|
||||
)
|
||||
store.save_sample(risky)
|
||||
|
||||
clean = SampleState(sample_id="bbb2", filename="clean.sys", source_path="C:/drivers/clean.sys")
|
||||
clean.mark_stage_completed("discover", data={"priority_score": 1.0})
|
||||
store.save_sample(clean)
|
||||
|
||||
d = store.sample_dir("aaa1")
|
||||
(d / "decompiled").mkdir(parents=True, exist_ok=True)
|
||||
(d / "decompiled" / "ghidra_result.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"success": True,
|
||||
"device_name": "RiskyDev",
|
||||
"symbolic_link": "\\\\DosDevices\\\\RiskyDev",
|
||||
"ioctl_handlers": [{"code": 0x222004}, {"code": 0x222008}],
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
if with_findings:
|
||||
(d / "findings.json").write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"rule_id": "pipelines.loldrivers.rules.ghidra-mmmapiospace-user-controlled",
|
||||
"severity": "HIGH",
|
||||
"message": "user controlled physical map",
|
||||
"line_start": 42,
|
||||
"matched_code": "MmMapIoSpace(pa, len, 0);",
|
||||
},
|
||||
{
|
||||
"rule_id": "pipelines.loldrivers.rules.method-neither",
|
||||
"severity": "MEDIUM",
|
||||
"message": "METHOD_NEITHER buffer",
|
||||
"line_start": 88,
|
||||
"matched_code": "irp->UserBuffer",
|
||||
},
|
||||
]
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
if with_assessment:
|
||||
(d / "assessment.md").write_text(
|
||||
"[VULNERABLE] arbitrary physical memory map via IOCTL 0x222004", encoding="utf-8"
|
||||
)
|
||||
return work
|
||||
|
||||
|
||||
class TestCollect:
|
||||
def test_totals(self, tmp_path):
|
||||
payload = collect(_seed(tmp_path))
|
||||
t = payload["totals"]
|
||||
assert t["samples"] == 2
|
||||
assert t["total_findings"] == 2
|
||||
assert t["with_findings"] == 1
|
||||
assert t["assessed"] == 1
|
||||
# severities are counted by whatever labels the findings actually use
|
||||
assert payload["severity_totals"] == {"HIGH": 1, "MEDIUM": 1}
|
||||
|
||||
def test_driver_detail_is_gathered(self, tmp_path):
|
||||
d = collect(_seed(tmp_path))["items"][0]
|
||||
assert d.name == "risky.sys"
|
||||
assert d.data["decompile.device_name"] == "RiskyDev"
|
||||
assert d.data["decompile.ioctl_count"] == 2
|
||||
assert d.data["discover.dangerous_imports"] == ["MmMapIoSpace"]
|
||||
assert "VULNERABLE" in d.texts["llm_output"]
|
||||
assert d.severity_counts == {"HIGH": 1, "MEDIUM": 1}
|
||||
assert "ghidra-mmmapiospace-user-controlled" in d.rule_hits
|
||||
|
||||
def test_only_interesting_drivers_are_kept_in_memory(self, tmp_path):
|
||||
payload = collect(_seed(tmp_path))
|
||||
# the clean driver is still counted and written to csv, but needs no page
|
||||
assert [d.sample_id for d in payload["items"]] == ["aaa1"]
|
||||
assert len(payload["rows"]) == 2
|
||||
|
||||
def test_bucketing_uses_the_assessment_verdict(self, tmp_path):
|
||||
payload = collect(_seed(tmp_path, classification="vulnerable"))
|
||||
assert payload["items"][0].bucket == BUCKET_VULNERABLE
|
||||
assert payload["buckets"][BUCKET_VULNERABLE] == 1
|
||||
|
||||
def test_findings_without_assessment_are_only_suspicious(self, tmp_path):
|
||||
payload = collect(_seed(tmp_path, with_assessment=False))
|
||||
assert payload["items"][0].bucket == BUCKET_SUSPICIOUS
|
||||
|
||||
def test_safe_verdict_is_not_flagged(self, tmp_path):
|
||||
payload = collect(_seed(tmp_path, classification="safe"))
|
||||
assert payload["buckets"].get(BUCKET_VULNERABLE, 0) == 0
|
||||
|
||||
def test_rule_totals_aggregated(self, tmp_path):
|
||||
rt = collect(_seed(tmp_path))["rule_totals"]
|
||||
assert rt["ghidra-mmmapiospace-user-controlled"] == 1
|
||||
assert rt["method-neither"] == 1
|
||||
|
||||
def test_run_metadata(self, tmp_path):
|
||||
run = collect(_seed(tmp_path))["run"]
|
||||
assert run["model"] == "claude-code/opus"
|
||||
assert run["pipeline"] == "loldrivers"
|
||||
|
||||
def test_safe_before_any_findings_exist(self, tmp_path):
|
||||
payload = collect(_seed(tmp_path, with_findings=False, with_assessment=False))
|
||||
assert payload["totals"]["total_findings"] == 0
|
||||
assert payload["totals"]["samples"] == 2
|
||||
|
||||
def test_empty_work_dir_is_safe(self, tmp_path):
|
||||
assert collect(tmp_path / "nothing")["totals"]["samples"] == 0
|
||||
|
||||
|
||||
class TestIndex:
|
||||
def test_self_contained_and_offline(self, tmp_path):
|
||||
out = render_index(collect(_seed(tmp_path)), tmp_path)
|
||||
assert out.startswith("<!doctype html>")
|
||||
# must open with no network: no remote assets of any kind
|
||||
assert "http://" not in out and "https://" not in out
|
||||
assert "<script src" not in out and "stylesheet" not in out
|
||||
|
||||
def test_vulnerable_is_the_headline(self, tmp_path):
|
||||
out = render_index(collect(_seed(tmp_path)), tmp_path)
|
||||
assert "assessed as vulnerable" in out
|
||||
# the verdict section leads, ahead of the supporting rule breakdown
|
||||
assert out.index("Vulnerable") < out.index("Rules that fired")
|
||||
assert "items/aaa1.html" in out
|
||||
|
||||
def test_empty_buckets_are_omitted(self, tmp_path):
|
||||
# nothing suspicious in this run, so no empty "Needs review" table
|
||||
out = render_index(collect(_seed(tmp_path)), tmp_path)
|
||||
assert "Needs review</h2>" not in out
|
||||
|
||||
def test_suspicious_section_appears_after_vulnerable(self, tmp_path):
|
||||
work = _seed(tmp_path)
|
||||
# add a second driver with findings but no assessment -> suspicious
|
||||
store = StateStore(work)
|
||||
s = SampleState(sample_id="ccc3", filename="maybe.sys", source_path="C:/drivers/maybe.sys")
|
||||
s.mark_stage_completed("semgrep_scanner", data={"finding_count": 1})
|
||||
store.save_sample(s)
|
||||
(store.sample_dir("ccc3") / "findings.json").write_text(
|
||||
json.dumps([{"rule_id": "r.x", "severity": "MEDIUM", "message": "m", "line_start": 1}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
out = render_index(collect(work), tmp_path)
|
||||
assert out.index("Vulnerable") < out.index("Needs review")
|
||||
|
||||
|
||||
class TestWriteReport:
|
||||
def test_writes_the_layered_output(self, tmp_path):
|
||||
index, json_path = write_report(_seed(tmp_path))
|
||||
assert index.name == "index.html" and index.parent.name == "report"
|
||||
out = index.parent
|
||||
for name in ("inventory.csv", "findings.jsonl", "report.json"):
|
||||
assert (out / name).exists(), name
|
||||
summary = json.loads(json_path.read_text(encoding="utf-8"))
|
||||
assert summary["totals"]["samples"] == 2
|
||||
# totals only - the summary must not grow with the corpus
|
||||
assert "items" not in summary
|
||||
|
||||
def test_index_stays_small_while_csv_holds_everything(self, tmp_path):
|
||||
index, _ = write_report(_seed(tmp_path))
|
||||
out = index.parent
|
||||
csv_text = (out / "inventory.csv").read_text(encoding="utf-8")
|
||||
assert "risky.sys" in csv_text and "clean.sys" in csv_text
|
||||
pages = {p.stem for p in (out / "items").glob("*.html")}
|
||||
assert pages == {"aaa1"}
|
||||
|
||||
def test_findings_are_one_json_object_per_line(self, tmp_path):
|
||||
index, _ = write_report(_seed(tmp_path))
|
||||
lines = [
|
||||
ln
|
||||
for ln in (index.parent / "findings.jsonl").read_text(encoding="utf-8").splitlines()
|
||||
if ln
|
||||
]
|
||||
assert len(lines) == 2
|
||||
rec = json.loads(lines[0])
|
||||
assert rec["sample_id"] == "aaa1" and rec["bucket"] == BUCKET_VULNERABLE
|
||||
assert rec["severity"] in ("HIGH", "MEDIUM")
|
||||
|
||||
def test_driver_page_shows_evidence_and_links_to_artifacts(self, tmp_path):
|
||||
index, _ = write_report(_seed(tmp_path))
|
||||
page = (index.parent / "items" / "aaa1.html").read_text(encoding="utf-8")
|
||||
assert "MmMapIoSpace(pa, len, 0);" in page
|
||||
assert "llm_output" in page and "ghidra_result" in page
|
||||
assert "all results" in page
|
||||
# links resolve to the real artifact folder
|
||||
assert "samples" in page
|
||||
|
||||
def test_untrusted_text_is_escaped_on_driver_pages(self, tmp_path):
|
||||
work = _seed(tmp_path)
|
||||
d = StateStore(work).sample_dir("aaa1")
|
||||
(d / "assessment.md").write_text("<img src=x onerror=alert(1)>", encoding="utf-8")
|
||||
index, _ = write_report(work)
|
||||
page = (index.parent / "items" / "aaa1.html").read_text(encoding="utf-8")
|
||||
# decompiled and LLM text is untrusted - never render it as live markup
|
||||
assert "<img src=x onerror" not in page
|
||||
assert "<img" in page
|
||||
|
||||
def test_capping_detail_is_disclosed_not_silent(self, tmp_path):
|
||||
index, _ = write_report(_seed(tmp_path), detail_limit=0)
|
||||
text = index.read_text(encoding="utf-8")
|
||||
assert "appear in" in text and "inventory.csv" in text
|
||||
|
||||
def test_custom_out_dir(self, tmp_path):
|
||||
out = tmp_path / "elsewhere"
|
||||
index, _ = write_report(_seed(tmp_path), out)
|
||||
assert index.parent == out
|
||||
|
||||
def test_regenerating_overwrites(self, tmp_path):
|
||||
work = _seed(tmp_path)
|
||||
write_report(work)
|
||||
index, _ = write_report(work)
|
||||
assert index.read_text(encoding="utf-8").count("<!doctype html>") == 1
|
||||
|
||||
|
||||
class TestPipelineAgnostic:
|
||||
"""a completely different pipeline must render without code changes."""
|
||||
|
||||
def _github_run(self, tmp_path):
|
||||
work = tmp_path / "work" / "srchunt"
|
||||
store = StateStore(work)
|
||||
store.save_run(
|
||||
RunState(
|
||||
run_id="run_9",
|
||||
pipeline="srchunt",
|
||||
target="github.com/acme",
|
||||
model="claude-code/opus",
|
||||
status=RunStatus.COMPLETED,
|
||||
)
|
||||
)
|
||||
repo = SampleState(sample_id="r1", filename="acme/api", source_path="github.com/acme/api")
|
||||
repo.mark_stage_completed("discover", data={"stars": 4200, "language": "go"})
|
||||
repo.mark_stage_completed(
|
||||
"sast", data={"finding_count": 1}, artifacts={"sast": "sast.json"}
|
||||
)
|
||||
repo.mark_stage_completed("triage", data={"verdict": "exploitable"})
|
||||
store.save_sample(repo)
|
||||
(store.sample_dir("r1") / "sast.json").write_text(
|
||||
json.dumps(
|
||||
[
|
||||
{
|
||||
"id": "go.sql-injection",
|
||||
"level": "critical",
|
||||
"title": "SQL injection in query builder",
|
||||
"path": "internal/db/query.go",
|
||||
"start_line": 88,
|
||||
"snippet": 'db.Raw("SELECT " + userInput)',
|
||||
}
|
||||
]
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return work
|
||||
|
||||
def test_declared_config_shapes_the_report(self, tmp_path):
|
||||
work = self._github_run(tmp_path)
|
||||
cfg = {
|
||||
"title": "Acme source review",
|
||||
"entity": "repository",
|
||||
"classification_key": "verdict",
|
||||
"vulnerable_when": ["exploitable"],
|
||||
"columns": ["discover.stars", "discover.language"],
|
||||
"findings_files": ["sast.json"],
|
||||
}
|
||||
index, json_path = write_report(work, config=cfg)
|
||||
text = index.read_text(encoding="utf-8")
|
||||
|
||||
# the pipeline's own vocabulary and verdict key drive the page
|
||||
assert "Acme source review" in text
|
||||
assert "repository" in text and "repositorys" not in text.replace("repositorys", "")
|
||||
assert "assessed as vulnerable" in text
|
||||
# its own severity vocabulary is preserved, not remapped to a fixed set
|
||||
assert "CRITICAL" in json.loads(json_path.read_text(encoding="utf-8"))["severity_totals"]
|
||||
# declared columns appear as chips/columns
|
||||
assert "stars" in text and "4200" in text
|
||||
|
||||
def test_works_with_no_config_at_all(self, tmp_path):
|
||||
work = self._github_run(tmp_path)
|
||||
index, _ = write_report(work)
|
||||
text = index.read_text(encoding="utf-8")
|
||||
# default entity wording, and the finding is still surfaced for review
|
||||
assert "sample" in text
|
||||
# no verdict key configured, so it lands in review rather than confirmed
|
||||
assert "Needs review" in text
|
||||
|
||||
def test_alien_finding_shape_is_normalized(self, tmp_path):
|
||||
work = self._github_run(tmp_path)
|
||||
payload = collect(work, config=ReportConfig.from_dict({"findings_files": ["sast.json"]}))
|
||||
f = payload["items"][0].findings[0]
|
||||
assert f["severity"] == "CRITICAL"
|
||||
assert f["rule_id"] == "go.sql-injection"
|
||||
assert "SQL injection" in f["message"]
|
||||
assert f["location"] == "internal/db/query.go"
|
||||
assert f["line"] == 88
|
||||
@@ -5,6 +5,7 @@ from pathlib import Path
|
||||
from deepzero.engine.runner import PipelineRunner
|
||||
from deepzero.engine.stage import (
|
||||
BulkMapProcessor,
|
||||
IngestProcessor,
|
||||
MapProcessor,
|
||||
ProcessorContext,
|
||||
ProcessorEntry,
|
||||
@@ -334,3 +335,87 @@ class TestResolveParallelism:
|
||||
proc = self._NoCeilingProc(StageSpec(name="d", processor="x"))
|
||||
spec = StageSpec(name="d", processor="x", parallel=0)
|
||||
assert PipelineRunner._resolve_parallelism(None, spec, proc) == 8
|
||||
|
||||
|
||||
class TestCachedSamplesStayInThePipeline:
|
||||
"""a sample whose work was already done still flows to the following
|
||||
stages. skipping work is not the same as filtering the sample out: a
|
||||
resumed run must analyse exactly what a fresh run would."""
|
||||
|
||||
class _AlreadyDone(MapProcessor):
|
||||
def should_skip(self, ctx, entry):
|
||||
return "output already cached"
|
||||
|
||||
def process(self, ctx, entry):
|
||||
raise AssertionError("process() must not run when should_skip returns a reason")
|
||||
|
||||
class _Counting(MapProcessor):
|
||||
seen: list = []
|
||||
|
||||
def process(self, ctx, entry):
|
||||
type(self).seen.append(entry.sample_id)
|
||||
return ProcessorResult.ok(data={"ran": True})
|
||||
|
||||
def test_cached_stage_is_completed_not_filtered(self, tmp_path):
|
||||
from deepzero.engine.state import SampleState
|
||||
from deepzero.engine.types import SampleStatus, StageStatus, Verdict
|
||||
|
||||
state = SampleState(sample_id="s1", filename="a.sys")
|
||||
state.mark_stage_cached("decompile", "output already cached")
|
||||
|
||||
out = state.history["decompile"]
|
||||
assert out.status == StageStatus.COMPLETED
|
||||
assert out.verdict == Verdict.CONTINUE
|
||||
assert out.skip_reason == "output already cached"
|
||||
# a skip reason is not an error
|
||||
assert out.error is None
|
||||
# and the sample is still eligible for later stages
|
||||
assert state.verdict != SampleStatus.FILTERED
|
||||
|
||||
def test_skip_reason_is_not_stored_as_an_error(self, tmp_path):
|
||||
from deepzero.engine.state import SampleState
|
||||
|
||||
state = SampleState(sample_id="s1", filename="a.sys")
|
||||
state.mark_stage_skipped("kernel_filter", "not a kernel driver")
|
||||
out = state.history["kernel_filter"]
|
||||
assert out.skip_reason == "not a kernel driver"
|
||||
assert out.error is None
|
||||
|
||||
def test_downstream_stage_still_sees_a_cached_sample(self, tmp_path):
|
||||
type(self._Counting).seen = []
|
||||
self._Counting.seen = []
|
||||
store = StateStore(tmp_path / "work")
|
||||
run = RunState(run_id="r1", pipeline="p")
|
||||
store.save_run(run)
|
||||
|
||||
target = tmp_path / "a.sys"
|
||||
target.write_bytes(b"MZ")
|
||||
|
||||
class OneSample(IngestProcessor):
|
||||
def __init__(self):
|
||||
self.spec = StageSpec(name="discover", processor="i")
|
||||
self.config = {}
|
||||
|
||||
def setup(self, global_config):
|
||||
pass
|
||||
|
||||
def process(self, ctx, t):
|
||||
return [Sample(sample_id="s1", source_path=target, filename="a.sys")]
|
||||
|
||||
cached_spec = StageSpec(name="decompile", processor="x")
|
||||
after_spec = StageSpec(name="scan", processor="y")
|
||||
after = self._Counting(after_spec)
|
||||
|
||||
runner = PipelineRunner(
|
||||
ingest=OneSample(),
|
||||
stages=[(cached_spec, self._AlreadyDone(cached_spec)), (after_spec, after)],
|
||||
state_store=store,
|
||||
pipeline_dir=tmp_path,
|
||||
global_config={},
|
||||
)
|
||||
runner.run(tmp_path, run)
|
||||
|
||||
# the whole point: the cached sample reached the next stage
|
||||
assert after.seen == ["s1"], "cached sample was dropped before the next stage"
|
||||
final = store.load_sample("s1")
|
||||
assert "scan" in final.history
|
||||
|
||||
@@ -24,8 +24,8 @@ class TestRulesPathResolution:
|
||||
)
|
||||
|
||||
def test_validate_and_process_resolve_the_same_path(self, tmp_path, monkeypatch):
|
||||
# regression: validate() approved a cwd-relative path while process()
|
||||
# used a pipeline_dir-relative one, so the scan ran with no rules.
|
||||
# the check that approves a run and the run itself must agree on where
|
||||
# the rules live, otherwise a scan can start with no rules loaded.
|
||||
# pin cwd so the cwd-relative "rules" resolves deterministically.
|
||||
monkeypatch.chdir(tmp_path)
|
||||
rules = tmp_path / "rules"
|
||||
|
||||
Reference in New Issue
Block a user