From 8fc5a996a04384ed61256f4092dfc05209a8bac7 Mon Sep 17 00:00:00 2001 From: Snehanjan Chatterjee Date: Fri, 20 Mar 2026 18:28:28 +0530 Subject: [PATCH] feat: add Risk Scenario Generator agent with CRML support Add a new ZAK agent (risk_scenario domain) that generates CRML cyber risk scenarios from domain intelligence. The agent fetches company homepage data, analyzes industry/tech stack, and produces schema-valid CRML YAML scenarios using FAIR-style (poisson+lognormal) and QBER-style (hierarchical_gamma_poisson +mixture) risk models. New files: - zak/agents/risk_scenario/ (agent, tools, __init__) - agent_templates/risk_scenario.yaml - tests/unit/test_risk_scenario.py (43 tests) Also fixes all ruff lint errors and reduces mypy strict-mode errors from 116 to 0 across the entire codebase. Co-Authored-By: Claude Opus 4.6 --- agent_templates/risk_scenario.yaml | 55 ++ tests/unit/test_risk_scenario.py | 871 +++++++++++++++++++ zak/agents/__init__.py | 1 + zak/agents/appsec/agent.py | 8 +- zak/agents/compliance/compliance_tools.py | 16 +- zak/agents/compliance/dpdp_agent.py | 4 +- zak/agents/compliance/iso27001_soc2_agent.py | 4 +- zak/agents/dep_patch/agent.py | 21 +- zak/agents/dep_patch/tools.py | 2 +- zak/agents/risk_quant/agent.py | 29 +- zak/agents/risk_scenario/__init__.py | 3 + zak/agents/risk_scenario/agent.py | 361 ++++++++ zak/agents/risk_scenario/tools.py | 463 ++++++++++ zak/agents/vuln_triage/agent.py | 10 +- zak/cli/main.py | 32 +- zak/cli/templates.py | 115 +++ zak/core/dsl/parser.py | 2 +- zak/core/dsl/schema.py | 15 +- zak/core/llm/anthropic_client.py | 6 +- zak/core/llm/local.py | 5 +- zak/core/llm/openai_client.py | 6 +- zak/core/runtime/llm_agent.py | 8 +- zak/core/runtime/registry.py | 6 +- zak/core/tools/builtins.py | 16 +- zak/core/tools/orchestration.py | 2 +- zak/core/tools/substrate.py | 12 +- zak/sif/graph/adapter.py | 4 +- zak/sif/telemetry/ingestor.py | 9 + 28 files changed, 1999 insertions(+), 87 deletions(-) create mode 100644 agent_templates/risk_scenario.yaml create mode 100644 tests/unit/test_risk_scenario.py create mode 100644 zak/agents/risk_scenario/__init__.py create mode 100644 zak/agents/risk_scenario/agent.py create mode 100644 zak/agents/risk_scenario/tools.py diff --git a/agent_templates/risk_scenario.yaml b/agent_templates/risk_scenario.yaml new file mode 100644 index 0000000..cfe4379 --- /dev/null +++ b/agent_templates/risk_scenario.yaml @@ -0,0 +1,55 @@ +agent: + id: risk-scenario-generator-v1 + name: "Risk Scenario Generator Agent" + domain: risk_scenario + version: "1.0.0" + +intent: + goal: "Generate calibrated CRML cyber risk scenarios for a target domain by gathering domain intelligence, identifying industry/threat profile, and producing schema-valid scenario YAML documents" + success_criteria: + - "At least 3 CRML scenarios generated" + - "All scenarios pass CRML schema validation" + - "Frequency and severity parameters are calibrated to industry/company size" + - "Each scenario includes relevant controls with effectiveness ratings" + priority: high + +reasoning: + mode: llm_react + autonomy_level: bounded + confidence_threshold: 0.75 + llm: + provider: openai + model: gpt-4o + temperature: 0.3 + max_iterations: 15 + max_tokens: 8192 + +capabilities: + tools: + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + data_access: + - web_public + graph_access: false + +boundaries: + risk_budget: low + allowed_actions: + - agent_execute + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + denied_actions: + - write_risk_node + - delete_asset + environment_scope: + - production + - staging + +safety: + guardrails: + - no_destructive_actions + - require_confidence_threshold + sandbox_profile: standard + audit_level: standard diff --git a/tests/unit/test_risk_scenario.py b/tests/unit/test_risk_scenario.py new file mode 100644 index 0000000..085a76a --- /dev/null +++ b/tests/unit/test_risk_scenario.py @@ -0,0 +1,871 @@ +""" +Tests for the Risk Scenario Generator Agent — tools, validation, and agent execution. +""" + +import json + +import yaml +import pytest + +from zak.core.dsl.schema import AgentDSL +from zak.core.runtime.agent import AgentContext, AgentResult +from zak.core.tools.substrate import ToolRegistry + +# Import tools to trigger @zak_tool registration +import zak.agents.risk_scenario.tools as scenario_tools # noqa: F401 +from zak.agents.risk_scenario.tools import ( + generate_crml_scenario, + validate_crml_scenario, +) +from zak.agents.risk_scenario.agent import RiskScenarioAgent + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +RISK_SCENARIO_YAML = """ +agent: + id: risk-scenario-test + name: "Risk Scenario Test Agent" + domain: risk_scenario + version: "1.0.0" + +intent: + goal: "Generate CRML risk scenarios for a target domain" + +reasoning: + mode: deterministic + autonomy_level: bounded + +capabilities: + tools: + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + +boundaries: + risk_budget: low + allowed_actions: + - agent_execute + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + +safety: + sandbox_profile: standard + audit_level: standard +""" + + +def make_context( + yaml_str: str = RISK_SCENARIO_YAML, + metadata: dict | None = None, +) -> AgentContext: + dsl = AgentDSL.model_validate(yaml.safe_load(yaml_str)) + return AgentContext( + tenant_id="test-tenant", + trace_id="trace-riskscen-001", + dsl=dsl, + environment="staging", + metadata=metadata or {"target_domain": "example.com"}, + ) + + +# --------------------------------------------------------------------------- +# Tests: Tool Registration +# --------------------------------------------------------------------------- + + +class TestToolRegistration: + def test_fetch_domain_intel_registered(self) -> None: + assert ToolRegistry.get().is_registered("fetch_domain_intel") + + def test_generate_crml_scenario_registered(self) -> None: + assert ToolRegistry.get().is_registered("generate_crml_scenario") + + def test_validate_crml_scenario_registered(self) -> None: + assert ToolRegistry.get().is_registered("validate_crml_scenario") + + +# --------------------------------------------------------------------------- +# Tests: generate_crml_scenario tool +# --------------------------------------------------------------------------- + + +class TestGenerateCRMLScenario: + def setup_method(self): + self.ctx = make_context() + + def test_basic_poisson_lognormal(self) -> None: + result = generate_crml_scenario( + context=self.ctx, + name="test-data-breach", + description="Test data breach scenario", + frequency_model="poisson", + frequency_lambda=0.05, + severity_model="lognormal", + severity_median="100 000", + severity_sigma=1.2, + ) + assert "yaml" in result + assert result["scenario_name"] == "test-data-breach" + + doc = yaml.safe_load(result["yaml"]) + assert doc["crml_scenario"] == "1.0" + assert doc["meta"]["name"] == "test-data-breach" + assert doc["scenario"]["frequency"]["model"] == "poisson" + assert doc["scenario"]["frequency"]["parameters"]["lambda"] == 0.05 + assert doc["scenario"]["severity"]["model"] == "lognormal" + assert doc["scenario"]["severity"]["parameters"]["sigma"] == 1.2 + + def test_hierarchical_gamma_poisson(self) -> None: + result = generate_crml_scenario( + context=self.ctx, + name="test-qber-model", + description="QBER-style Bayesian model", + frequency_model="hierarchical_gamma_poisson", + frequency_lambda=0.0, # not used for this model + frequency_alpha_base=2.0, + frequency_beta_base=1.5, + severity_model="lognormal", + severity_median="250 000", + severity_sigma=1.5, + ) + doc = yaml.safe_load(result["yaml"]) + freq = doc["scenario"]["frequency"] + assert freq["model"] == "hierarchical_gamma_poisson" + assert freq["parameters"]["alpha_base"] == 2.0 + assert freq["parameters"]["beta_base"] == 1.5 + + def test_mixture_severity(self) -> None: + components = [ + {"lognormal": {"weight": 0.7, "median": "200 000", "sigma": 1.2, "currency": "USD"}}, + {"gamma": {"weight": 0.3, "shape": 2.5, "scale": "10 000", "currency": "USD"}}, + ] + result = generate_crml_scenario( + context=self.ctx, + name="test-mixture", + description="Mixture severity model", + frequency_model="poisson", + frequency_lambda=0.1, + severity_model="mixture", + severity_median="200 000", + severity_sigma=1.2, + severity_components_json=json.dumps(components), + ) + doc = yaml.safe_load(result["yaml"]) + assert doc["scenario"]["severity"]["model"] == "mixture" + assert len(doc["scenario"]["severity"]["components"]) == 2 + + def test_with_controls(self) -> None: + controls = [ + {"id": "org:iam.mfa", "effectiveness_against_threat": 0.85}, + {"id": "org:email.dmarc", "effectiveness_against_threat": 0.55}, + ] + result = generate_crml_scenario( + context=self.ctx, + name="test-with-controls", + description="Scenario with controls", + frequency_model="poisson", + frequency_lambda=0.4, + severity_model="lognormal", + severity_median="25 000", + severity_sigma=1.15, + controls_json=json.dumps(controls), + ) + doc = yaml.safe_load(result["yaml"]) + assert len(doc["scenario"]["controls"]) == 2 + assert doc["scenario"]["controls"][0]["id"] == "org:iam.mfa" + assert doc["scenario"]["controls"][0]["effectiveness_against_threat"] == 0.85 + + def test_tags_and_metadata(self) -> None: + result = generate_crml_scenario( + context=self.ctx, + name="test-meta", + description="Test metadata", + frequency_model="poisson", + frequency_lambda=0.1, + severity_model="lognormal", + severity_median="50 000", + severity_sigma=1.0, + tags="ransomware,encryption", + company_size="enterprise,large-enterprise", + industries="healthcare,finance", + author="Test Author", + ) + doc = yaml.safe_load(result["yaml"]) + assert doc["meta"]["tags"] == ["ransomware", "encryption"] + assert doc["meta"]["company_size"] == ["enterprise", "large-enterprise"] + assert doc["meta"]["industries"] == ["healthcare", "finance"] + assert doc["meta"]["author"] == "Test Author" + + def test_per_asset_unit_basis(self) -> None: + result = generate_crml_scenario( + context=self.ctx, + name="test-per-asset", + description="Per-asset-unit basis", + frequency_model="poisson", + frequency_lambda=0.01, + severity_model="lognormal", + severity_median="50 000", + severity_sigma=1.0, + frequency_basis="per_asset_unit_per_year", + ) + doc = yaml.safe_load(result["yaml"]) + assert doc["scenario"]["frequency"]["basis"] == "per_asset_unit_per_year" + + def test_default_mixture_components_when_json_invalid(self) -> None: + result = generate_crml_scenario( + context=self.ctx, + name="test-bad-mixture-json", + description="Bad JSON falls back to defaults", + frequency_model="poisson", + frequency_lambda=0.1, + severity_model="mixture", + severity_median="100 000", + severity_sigma=1.0, + severity_components_json="NOT VALID JSON", + ) + doc = yaml.safe_load(result["yaml"]) + assert doc["scenario"]["severity"]["model"] == "mixture" + # Should fall back to default components + assert len(doc["scenario"]["severity"]["components"]) == 2 + + def test_gamma_severity_model(self) -> None: + result = generate_crml_scenario( + context=self.ctx, + name="test-gamma-severity", + description="Gamma severity model", + frequency_model="poisson", + frequency_lambda=0.05, + severity_model="gamma", + severity_median="50 000", + severity_sigma=2.5, + ) + doc = yaml.safe_load(result["yaml"]) + sev = doc["scenario"]["severity"] + assert sev["model"] == "gamma" + assert sev["parameters"]["shape"] == 2.5 + assert sev["parameters"]["scale"] == "50 000" + + +# --------------------------------------------------------------------------- +# Tests: validate_crml_scenario tool +# --------------------------------------------------------------------------- + + +class TestValidateCRMLScenario: + def setup_method(self): + self.ctx = make_context() + + def test_valid_simple_scenario(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test-scenario" +scenario: + frequency: + basis: per_organization_per_year + model: poisson + parameters: + lambda: 0.05 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.2 + currency: USD +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is True + assert len(result["errors"]) == 0 + + def test_missing_crml_scenario_version(self) -> None: + yaml_str = """ +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.05 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("crml_scenario" in e for e in result["errors"]) + + def test_missing_scenario_block(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("scenario" in e.lower() for e in result["errors"]) + + def test_missing_meta_name(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: {} +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: lognormal + parameters: + median: "50 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("meta.name" in e for e in result["errors"]) + + def test_invalid_frequency_model(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: exponential + parameters: + rate: 0.1 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("frequency model" in e.lower() for e in result["errors"]) + + def test_missing_poisson_lambda(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: {} + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("lambda" in e.lower() for e in result["errors"]) + + def test_negative_lambda_rejected(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: -0.5 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("non-negative" in e.lower() for e in result["errors"]) + + def test_high_lambda_warning(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 150 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is True # Warning, not error + assert any("unusually high" in w.lower() for w in result["warnings"]) + + def test_missing_lognormal_sigma(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: lognormal + parameters: + median: "100 000" +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("sigma" in e.lower() for e in result["errors"]) + + def test_missing_lognormal_median_and_mu(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: lognormal + parameters: + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("median" in e.lower() or "mu" in e.lower() for e in result["errors"]) + + def test_invalid_severity_model(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: exponential + parameters: + rate: 0.01 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("severity model" in e.lower() for e in result["errors"]) + + def test_invalid_control_id_attck_prefix(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 + controls: + - id: "attck:T1566" + effectiveness_against_threat: 0.5 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("attck:" in e for e in result["errors"]) + + def test_control_effectiveness_out_of_range(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 + controls: + - id: "org:iam.mfa" + effectiveness_against_threat: 1.5 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("0.0-1.0" in e for e in result["errors"]) + + def test_mixture_empty_components(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: mixture + parameters: {} +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("component" in e.lower() for e in result["errors"]) + + def test_mixture_weights_warning(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: mixture + parameters: {} + components: + - lognormal: + weight: 0.5 + median: "100 000" + sigma: 1.0 + - gamma: + weight: 0.3 + shape: 2.0 + scale: 5000 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is True # Valid but warning about weights + assert any("weights sum" in w.lower() for w in result["warnings"]) + + def test_hierarchical_missing_alpha(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: hierarchical_gamma_poisson + parameters: + beta_base: 1.5 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("alpha_base" in e for e in result["errors"]) + + def test_invalid_yaml(self) -> None: + result = validate_crml_scenario(context=self.ctx, yaml_content=": : invalid [yaml") + assert result["valid"] is False + assert any("parse error" in e.lower() for e in result["errors"]) + + def test_invalid_frequency_basis(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + basis: per_employee_per_year + model: poisson + parameters: + lambda: 0.1 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: 1.0 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("basis" in e.lower() for e in result["errors"]) + + def test_using_mu_instead_of_median(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.8 + severity: + model: lognormal + parameters: + mu: 10.0 + sigma: 1.2 + currency: USD +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is True + + def test_negative_sigma_rejected(self) -> None: + yaml_str = """ +crml_scenario: "1.0" +meta: + name: "test" +scenario: + frequency: + model: poisson + parameters: + lambda: 0.1 + severity: + model: lognormal + parameters: + median: "100 000" + sigma: -0.5 +""" + result = validate_crml_scenario(context=self.ctx, yaml_content=yaml_str) + assert result["valid"] is False + assert any("sigma" in e.lower() for e in result["errors"]) + + +# --------------------------------------------------------------------------- +# Tests: Roundtrip — generate then validate +# --------------------------------------------------------------------------- + + +class TestRoundtrip: + def setup_method(self): + self.ctx = make_context() + + def test_generated_poisson_lognormal_validates(self) -> None: + result = generate_crml_scenario( + context=self.ctx, + name="roundtrip-test", + description="Roundtrip validation test", + frequency_model="poisson", + frequency_lambda=0.1, + severity_model="lognormal", + severity_median="250 000", + severity_sigma=1.5, + tags="test,roundtrip", + ) + validation = validate_crml_scenario(context=self.ctx, yaml_content=result["yaml"]) + assert validation["valid"] is True, f"Validation failed: {validation['errors']}" + + def test_generated_hierarchical_mixture_validates(self) -> None: + components = [ + {"lognormal": {"weight": 0.7, "median": "200 000", "sigma": 1.2, "currency": "USD"}}, + {"gamma": {"weight": 0.3, "shape": 2.5, "scale": "10 000", "currency": "USD"}}, + ] + result = generate_crml_scenario( + context=self.ctx, + name="roundtrip-qber", + description="QBER roundtrip test", + frequency_model="hierarchical_gamma_poisson", + frequency_lambda=0.0, + frequency_alpha_base=2.0, + frequency_beta_base=1.5, + severity_model="mixture", + severity_median="200 000", + severity_sigma=1.2, + severity_components_json=json.dumps(components), + ) + validation = validate_crml_scenario(context=self.ctx, yaml_content=result["yaml"]) + assert validation["valid"] is True, f"Validation failed: {validation['errors']}" + + def test_generated_with_controls_validates(self) -> None: + controls = [ + {"id": "org:iam.mfa", "effectiveness_against_threat": 0.85}, + {"id": "org:net.firewall", "effectiveness_against_threat": 0.6}, + {"id": "org:email.dmarc", "effectiveness_against_threat": 0.55}, + ] + result = generate_crml_scenario( + context=self.ctx, + name="roundtrip-controls", + description="Roundtrip with controls", + frequency_model="poisson", + frequency_lambda=0.4, + severity_model="lognormal", + severity_median="25 000", + severity_sigma=1.15, + controls_json=json.dumps(controls), + ) + validation = validate_crml_scenario(context=self.ctx, yaml_content=result["yaml"]) + assert validation["valid"] is True, f"Validation failed: {validation['errors']}" + + +# --------------------------------------------------------------------------- +# Tests: Agent Registration +# --------------------------------------------------------------------------- + + +class TestAgentRegistration: + def test_risk_scenario_agent_importable(self) -> None: + from zak.agents.risk_scenario.agent import RiskScenarioAgent + assert RiskScenarioAgent is not None + + def test_domain_attribute(self) -> None: + assert RiskScenarioAgent._zak_domain == "risk_scenario" + + def test_edition_open_source(self) -> None: + assert RiskScenarioAgent._zak_edition == "open-source" + + def test_version(self) -> None: + assert RiskScenarioAgent._zak_version == "1.0.0" + + +# --------------------------------------------------------------------------- +# Tests: DSL Validation +# --------------------------------------------------------------------------- + + +class TestDSLValidation: + def test_risk_scenario_yaml_validates(self) -> None: + yaml_str = """ +agent: + id: risk-scenario-gen-v1 + name: "Risk Scenario Generator Agent" + domain: risk_scenario + version: "1.0.0" + +intent: + goal: "Generate CRML risk scenarios" + +reasoning: + mode: llm_react + autonomy_level: bounded + confidence_threshold: 0.75 + llm: + provider: openai + model: gpt-4o + temperature: 0.3 + max_iterations: 15 + +capabilities: + tools: + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + +boundaries: + risk_budget: low + allowed_actions: + - agent_execute + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + +safety: + sandbox_profile: standard + audit_level: standard +""" + dsl = AgentDSL.model_validate(yaml.safe_load(yaml_str)) + assert dsl.agent.domain == "risk_scenario" + assert dsl.reasoning.mode.value == "llm_react" + + def test_deterministic_mode_validates(self) -> None: + dsl = AgentDSL.model_validate(yaml.safe_load(RISK_SCENARIO_YAML)) + assert dsl.agent.domain == "risk_scenario" + assert dsl.reasoning.mode.value == "deterministic" + + +# --------------------------------------------------------------------------- +# Tests: CRML Spec Compliance +# --------------------------------------------------------------------------- + + +class TestCRMLSpecCompliance: + """Verify generated scenarios conform to actual CRML spec patterns.""" + + def setup_method(self): + self.ctx = make_context() + + def test_matches_crml_data_breach_simple_pattern(self) -> None: + """Generated scenario should match the structure of examples/scenarios/data-breach-simple.yaml.""" + result = generate_crml_scenario( + context=self.ctx, + name="data-breach-simple", + description="Simple data breach risk model for beginners", + frequency_model="poisson", + frequency_lambda=0.05, + severity_model="lognormal", + severity_median="100 000", + severity_sigma=1.2, + severity_currency="USD", + tags="data-breach,beginner,pii", + company_size="smb", + industries="all", + ) + doc = yaml.safe_load(result["yaml"]) + + # Structural match with CRML examples + assert doc["crml_scenario"] == "1.0" + assert "meta" in doc + assert "scenario" in doc + assert "frequency" in doc["scenario"] + assert "severity" in doc["scenario"] + assert doc["scenario"]["frequency"]["basis"] == "per_organization_per_year" + assert doc["scenario"]["severity"]["parameters"]["currency"] == "USD" + + def test_matches_crml_ransomware_pattern(self) -> None: + """Generated scenario should match ransomware-scenario.yaml patterns.""" + result = generate_crml_scenario( + context=self.ctx, + name="ransomware-scenario", + description="Real-world ransomware risk model based on industry statistics", + frequency_model="poisson", + frequency_lambda=0.08, + severity_model="lognormal", + severity_median="700 000", + severity_sigma=1.8, + severity_currency="USD", + tags="ransomware,extortion", + company_size="enterprise,large-enterprise", + ) + doc = yaml.safe_load(result["yaml"]) + + assert doc["scenario"]["frequency"]["parameters"]["lambda"] == 0.08 + assert doc["scenario"]["severity"]["parameters"]["median"] == "700 000" + assert doc["scenario"]["severity"]["parameters"]["sigma"] == 1.8 + + def test_matches_crml_qber_simplified_pattern(self) -> None: + """Generated QBER scenario should match qber-simplified.yaml patterns.""" + components = [ + {"lognormal": {"weight": 0.7, "median": "162 755", "currency": "USD", "sigma": 1.2}}, + {"gamma": {"weight": 0.3, "shape": 2.5, "scale": 10000, "currency": "USD"}}, + ] + result = generate_crml_scenario( + context=self.ctx, + name="qber-simplified", + description="Simplified QBER-style model", + frequency_model="hierarchical_gamma_poisson", + frequency_lambda=0.0, + frequency_alpha_base=1.5, + frequency_beta_base=1.5, + severity_model="mixture", + severity_median="162 755", + severity_sigma=1.2, + severity_components_json=json.dumps(components), + tags="qber,bayesian", + company_size="enterprise,large-enterprise", + ) + doc = yaml.safe_load(result["yaml"]) + + assert doc["scenario"]["frequency"]["model"] == "hierarchical_gamma_poisson" + assert doc["scenario"]["severity"]["model"] == "mixture" + assert len(doc["scenario"]["severity"]["components"]) == 2 diff --git a/zak/agents/__init__.py b/zak/agents/__init__.py index 5916c08..48893db 100644 --- a/zak/agents/__init__.py +++ b/zak/agents/__init__.py @@ -11,6 +11,7 @@ "zak.agents.code_auditor.agent", "zak.agents.dep_patch.agent", "zak.agents.slopsquatting.agent", + "zak.agents.risk_scenario.agent", ] diff --git a/zak/agents/appsec/agent.py b/zak/agents/appsec/agent.py index fd228b4..1f0c9c9 100644 --- a/zak/agents/appsec/agent.py +++ b/zak/agents/appsec/agent.py @@ -51,15 +51,15 @@ def execute(self, context: AgentContext) -> AgentResult: def _execute_deterministic(self, context: AgentContext) -> AgentResult: tenant_id = context.tenant_id repos = ( - self._adapter.get_nodes(tenant_id=tenant_id, node_type="repository") # type: ignore[union-attr] + self._adapter.get_nodes(tenant_id=tenant_id, node_type="repository") # type: ignore[attr-defined] if self._adapter is not None else [] ) deps = ( - self._adapter.get_nodes(tenant_id=tenant_id, node_type="dependency") # type: ignore[union-attr] + self._adapter.get_nodes(tenant_id=tenant_id, node_type="dependency") # type: ignore[attr-defined] if self._adapter is not None else [] ) - findings: list[dict] = [] + findings: list[dict[str, object]] = [] secrets_found = 0 for repo in repos: @@ -100,7 +100,7 @@ class _LLMAppSecAgent(LLMAgent): """LLM-powered application security analysis using the ReAct loop.""" @property - def tools(self) -> list: + def tools(self) -> list[object]: from zak.core.tools.builtins import ( list_assets, list_vulnerabilities, diff --git a/zak/agents/compliance/compliance_tools.py b/zak/agents/compliance/compliance_tools.py index 5c62b82..68d0163 100644 --- a/zak/agents/compliance/compliance_tools.py +++ b/zak/agents/compliance/compliance_tools.py @@ -7,9 +7,9 @@ from __future__ import annotations -import json import os from datetime import date +from typing import Any from zak.core.runtime.agent import AgentContext from zak.core.tools.substrate import zak_tool @@ -18,7 +18,7 @@ # ISO 27001:2022 — 93 controls across 4 themes # --------------------------------------------------------------------------- -ISO27001_CONTROLS: list[dict] = [ +ISO27001_CONTROLS: list[dict[str, Any]] = [ # Theme 1: Organisational Controls (A.5) { "theme": "Organisational Controls", @@ -147,7 +147,7 @@ # SOC 2 Trust Service Criteria (2017 + 2022 updates) # --------------------------------------------------------------------------- -SOC2_CRITERIA: list[dict] = [ +SOC2_CRITERIA: list[dict[str, Any]] = [ { "category": "Common Criteria (CC) — Security", "id": "CC", @@ -270,7 +270,7 @@ action_id="get_iso27001_controls", tags=["compliance", "iso27001", "read"], ) -def get_iso27001_controls(context: AgentContext) -> dict: +def get_iso27001_controls(context: AgentContext) -> dict[str, Any]: """Return all ISO 27001:2022 Annex A controls grouped by theme.""" total = sum(len(theme["controls"]) for theme in ISO27001_CONTROLS) return { @@ -286,7 +286,7 @@ def get_iso27001_controls(context: AgentContext) -> dict: action_id="get_soc2_criteria", tags=["compliance", "soc2", "read"], ) -def get_soc2_criteria(context: AgentContext) -> dict: +def get_soc2_criteria(context: AgentContext) -> dict[str, Any]: """Return all SOC 2 Trust Service Criteria grouped by category.""" return { "framework": "SOC 2 (AICPA Trust Service Criteria 2017)", @@ -305,7 +305,7 @@ def save_policy_document( policy_name: str, policy_content: str, framework: str, -) -> dict: +) -> dict[str, Any]: """ Saves a draft policy document as a Markdown file. @@ -355,7 +355,7 @@ def save_policy_document( def save_gap_report( context: AgentContext, report_content: str, -) -> dict: +) -> dict[str, Any]: """ Saves the master compliance gap report as a Markdown file. @@ -397,7 +397,7 @@ def save_gap_report( action_id="list_output_files", tags=["compliance", "read"], ) -def list_output_files(context: AgentContext) -> dict: +def list_output_files(context: AgentContext) -> dict[str, Any]: """Lists all files saved to the compliance output directory.""" out_dir = context.metadata.get("output_dir", "/tmp/zak_compliance_output") if not os.path.exists(out_dir): diff --git a/zak/agents/compliance/dpdp_agent.py b/zak/agents/compliance/dpdp_agent.py index c3f09cf..4c00fc7 100644 --- a/zak/agents/compliance/dpdp_agent.py +++ b/zak/agents/compliance/dpdp_agent.py @@ -4,6 +4,8 @@ from __future__ import annotations +from typing import Any + from zak.core.runtime.agent import AgentContext from zak.core.runtime.llm_agent import LLMAgent from zak.core.runtime.registry import register_agent @@ -42,5 +44,5 @@ def system_prompt(self, context: AgentContext) -> str: """ @property - def tools(self) -> list: + def tools(self) -> list[Any]: return [website_tools.fetch_website_content] diff --git a/zak/agents/compliance/iso27001_soc2_agent.py b/zak/agents/compliance/iso27001_soc2_agent.py index 4cc5e1b..7a98b99 100644 --- a/zak/agents/compliance/iso27001_soc2_agent.py +++ b/zak/agents/compliance/iso27001_soc2_agent.py @@ -19,6 +19,8 @@ from __future__ import annotations +from typing import Any + from zak.core.runtime.agent import AgentContext from zak.core.runtime.llm_agent import LLMAgent from zak.core.runtime.registry import register_agent @@ -173,7 +175,7 @@ def system_prompt(self, context: AgentContext) -> str: """ @property - def tools(self) -> list: + def tools(self) -> list[Any]: return [ compliance_tools.get_iso27001_controls, compliance_tools.get_soc2_criteria, diff --git a/zak/agents/dep_patch/agent.py b/zak/agents/dep_patch/agent.py index 95da9bd..97e80dd 100644 --- a/zak/agents/dep_patch/agent.py +++ b/zak/agents/dep_patch/agent.py @@ -23,6 +23,7 @@ import json import os from datetime import datetime, timezone +from typing import Any from zak.core.dsl.schema import ReasoningMode from zak.core.runtime.agent import AgentContext, AgentResult, BaseAgent @@ -88,7 +89,7 @@ def _execute_deterministic(self, context: AgentContext) -> AgentResult: non_updatable_deps = [d for d in deps if not d["is_updatable"]] # 3. Check registry for updates - results: list[dict] = [] + results: list[dict[str, Any]] = [] # Non-updatable: pass through with empty versions for dep in non_updatable_deps: @@ -131,7 +132,7 @@ def _execute_deterministic(self, context: AgentContext) -> AgentResult: ] # 5. Risk assessment - assessed_updates: list[dict] = [] + assessed_updates: list[dict[str, Any]] = [] if updatable_results: assessed_updates = ToolExecutor.call( assess_update_risks, context=context, updates=updatable_results, @@ -143,7 +144,7 @@ def _execute_deterministic(self, context: AgentContext) -> AgentResult: results[i] = assessed_by_name[r["name"]] # 6. Build summary - summary: dict = { + summary: dict[str, Any] = { "total_dependencies": len(deps), "updatable": len([ r for r in results @@ -267,7 +268,7 @@ def execute(self, context: AgentContext) -> AgentResult: tools_schema = _build_openai_schema(available_tools) # LLM config from DSL - llm_cfg: dict = {} + llm_cfg: dict[str, Any] = {} if context.dsl.reasoning.llm: llm_block = context.dsl.reasoning.llm llm_cfg = ( @@ -309,7 +310,7 @@ def execute(self, context: AgentContext) -> AgentResult: "Ground every decision in tool output. Do not invent version numbers." ) - messages: list[dict] = [ + messages: list[dict[str, Any]] = [ {"role": "system", "content": system}, { "role": "user", @@ -320,12 +321,12 @@ def execute(self, context: AgentContext) -> AgentResult: }, ] - reasoning_trace: list[dict] = [] - total_usage: dict = { + reasoning_trace: list[dict[str, Any]] = [] + total_usage: dict[str, Any] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, } - def resolve(name: str): # type: ignore[return] + def resolve(name: str) -> Any: for fn in available_tools: meta = getattr(fn, "_zak_tool", None) if meta and meta.action_id == name: @@ -362,9 +363,9 @@ def resolve(name: str): # type: ignore[return] ) # Process tool calls - tool_results: list[dict] = [] + tool_results: list[dict[str, Any]] = [] for tc in response.tool_calls: - entry: dict = { + entry: dict[str, Any] = { "iteration": iteration + 1, "type": "tool_call", "tool": tc.name, diff --git a/zak/agents/dep_patch/tools.py b/zak/agents/dep_patch/tools.py index 496963f..ec7091a 100644 --- a/zak/agents/dep_patch/tools.py +++ b/zak/agents/dep_patch/tools.py @@ -182,7 +182,7 @@ def _github_api_request( req.add_header("Content-Type", "application/json") ctx = ssl.create_default_context() with urllib.request.urlopen(req, context=ctx) as resp: - return json.loads(resp.read().decode("utf-8")) + return json.loads(resp.read().decode("utf-8")) # type: ignore[no-any-return] # --------------------------------------------------------------------------- diff --git a/zak/agents/risk_quant/agent.py b/zak/agents/risk_quant/agent.py index e509e27..d57e526 100644 --- a/zak/agents/risk_quant/agent.py +++ b/zak/agents/risk_quant/agent.py @@ -21,6 +21,8 @@ from __future__ import annotations +from typing import Any + from zak.core.dsl.schema import ReasoningMode from zak.core.runtime.agent import AgentContext, AgentResult, BaseAgent from zak.core.runtime.registry import register_agent @@ -68,12 +70,12 @@ def _execute_llm(self, context: AgentContext) -> AgentResult: def _execute_deterministic(self, context: AgentContext) -> AgentResult: tenant_id = context.tenant_id - scored: list[dict] = [] + scored: list[dict[str, Any]] = [] errors: list[str] = [] # Load all assets for this tenant (graceful when graph is unavailable) assets = ( - self._adapter.get_nodes(tenant_id=tenant_id, node_type="asset") + self._adapter.get_nodes(tenant_id=tenant_id, node_type="asset") # type: ignore[attr-defined] if self._adapter is not None else [] ) @@ -88,9 +90,12 @@ def _execute_deterministic(self, context: AgentContext) -> AgentResult: risk_score=risk_output.risk_score, eal=self._compute_eal_stub(risk_output.risk_score), source=context.agent_id, + valid_to=None, + confidence=0.95, + var_95=None, ) if self._adapter is not None: - self._adapter.upsert_node(tenant_id, risk_node) + self._adapter.upsert_node(tenant_id, risk_node) # type: ignore[attr-defined] scored.append({ "asset_id": asset["node_id"], "risk_score": risk_output.risk_score, @@ -111,19 +116,19 @@ def _execute_deterministic(self, context: AgentContext) -> AgentResult: }, ) - def _score_asset(self, asset: dict, tenant_id: str) -> object: + def _score_asset(self, asset: dict[str, Any], tenant_id: str) -> Any: """Compute risk for a single asset dict.""" criticality = asset.get("criticality", "medium") exposure = asset.get("exposure_level", "internal") # Load worst-case vulnerability exploitability for this asset - vulns = self._adapter.get_nodes(tenant_id=tenant_id, node_type="vulnerability") if self._adapter is not None else [] + vulns = self._adapter.get_nodes(tenant_id=tenant_id, node_type="vulnerability") if self._adapter is not None else [] # type: ignore[attr-defined] max_exploitability = max( (float(v.get("exploitability", 0.5)) for v in vulns), default=0.5 ) # Load best control effectiveness - controls = self._adapter.get_nodes(tenant_id=tenant_id, node_type="control") if self._adapter is not None else [] + controls = self._adapter.get_nodes(tenant_id=tenant_id, node_type="control") if self._adapter is not None else [] # type: ignore[attr-defined] max_control_eff = max( (float(c.get("effectiveness", 0.5)) for c in controls), default=0.0 ) @@ -185,7 +190,7 @@ def execute(self, context: AgentContext) -> AgentResult: tools_schema = _build_openai_schema(available_tools) # LLM config from DSL - llm_cfg: dict = {} + llm_cfg: dict[str, Any] = {} if context.dsl.reasoning.llm: llm_block = context.dsl.reasoning.llm llm_cfg = ( @@ -230,10 +235,10 @@ def execute(self, context: AgentContext) -> AgentResult: }, ] - reasoning_trace = [] - total_usage: dict = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + reasoning_trace: list[Any] = [] + total_usage: dict[str, Any] = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} - def resolve(name: str): + def resolve(name: str) -> Any: for fn in available_tools: meta = getattr(fn, "_zak_tool", None) if meta and meta.action_id == name: @@ -300,14 +305,14 @@ def resolve(name: str): messages.append({ "role": "assistant", - "content": response.content, + "content": response.content or "", "tool_calls": [ { "id": tc.id, "type": "function", "function": {"name": tc.name, "arguments": json.dumps(tc.arguments)}, } for tc in response.tool_calls - ], + ], # type: ignore[dict-item] }) messages.extend(tool_results) diff --git a/zak/agents/risk_scenario/__init__.py b/zak/agents/risk_scenario/__init__.py new file mode 100644 index 0000000..36eed3c --- /dev/null +++ b/zak/agents/risk_scenario/__init__.py @@ -0,0 +1,3 @@ +"""Risk Scenario Generator Agent — generates CRML cyber risk scenarios from domain intelligence.""" + +from zak.agents.risk_scenario import tools as scenario_tools # noqa: F401 diff --git a/zak/agents/risk_scenario/agent.py b/zak/agents/risk_scenario/agent.py new file mode 100644 index 0000000..6755618 --- /dev/null +++ b/zak/agents/risk_scenario/agent.py @@ -0,0 +1,361 @@ +""" +Risk Scenario Generator Agent — ZAK reference implementation. + +Takes a domain as input and generates CRML (Cyber Risk Modeling Language) cyber +risk scenarios for the company associated with that domain. + +Uses the CRML specification (https://github.com/Faux16/crml) to produce +schema-valid scenario YAML documents with calibrated frequency and severity +parameters based on publicly available domain intelligence. + +Supports two CRML modeling styles: + - FAIR-style: Poisson frequency + Lognormal severity (simpler, expert-driven) + - QBER-style: Hierarchical Gamma-Poisson + Mixture severity (Bayesian, enterprise) + +Execution modes: + reasoning.mode: llm_react (recommended — LLM analyses domain and generates scenarios) + reasoning.mode: deterministic (generates a basic data-breach scenario from domain info) +""" + +from __future__ import annotations + +from typing import Any + +from zak.core.runtime.agent import AgentContext, AgentResult, BaseAgent +from zak.core.runtime.registry import register_agent +from zak.agents.risk_scenario.tools import ( + fetch_domain_intel, + generate_crml_scenario, + validate_crml_scenario, +) + + +@register_agent( + domain="risk_scenario", + description=( + "Generates CRML cyber risk scenarios for a target domain. " + "Performs domain reconnaissance, identifies industry and threat profile, " + "then produces calibrated CRML scenario YAML documents using FAIR or QBER models." + ), + version="1.0.0", + edition="open-source", +) +class RiskScenarioAgent(BaseAgent): + """ + Generates CRML cyber risk scenarios from domain intelligence. + + In LLM mode, the agent follows this tool sequence: + 1. fetch_domain_intel -> gather company/domain context + 2. generate_crml_scenario -> produce CRML YAML for each identified risk + 3. validate_crml_scenario -> verify CRML schema compliance + 4. STOP + summarize -> return scenarios + risk narrative + + In deterministic mode, produces a baseline data-breach scenario + using conservative industry-average parameters. + """ + + def execute(self, context: AgentContext) -> AgentResult: + from zak.core.dsl.schema import ReasoningMode + + if context.dsl.reasoning.mode == ReasoningMode.LLM_REACT: + return self._execute_llm(context) + return self._execute_deterministic(context) + + # ── LLM-powered execution ─────────────────────────────────────────── + + def _execute_llm(self, context: AgentContext) -> AgentResult: + """Full LLM ReAct loop for intelligent scenario generation.""" + agent = _LLMRiskScenarioAgent() + return agent.execute(context) + + # ── Deterministic fallback ────────────────────────────────────────── + + def _execute_deterministic(self, context: AgentContext) -> AgentResult: + """ + Generate a baseline CRML scenario using conservative defaults. + + Fetches domain intelligence, then produces a single data-breach + scenario with industry-average FAIR parameters. + """ + from zak.core.tools.substrate import ToolExecutor + + target_domain = context.metadata.get("target_domain", "") + if not target_domain: + return AgentResult.fail( + context, + errors=["No target_domain provided in agent metadata."], + ) + + # Step 1: Gather domain intel + try: + intel = ToolExecutor.call( + fetch_domain_intel, context=context, domain=target_domain + ) + except Exception as exc: + intel = {"domain": target_domain, "error": str(exc), "page_text": ""} + + domain_name = intel.get("domain", target_domain) + has_hsts = "strict-transport-security" in str(intel.get("http_headers", {})).lower() + + # Step 2: Generate a conservative data-breach scenario + scenario_result = ToolExecutor.call( + generate_crml_scenario, + context=context, + name=f"{domain_name}-data-breach", + description=f"Baseline data breach risk scenario for {domain_name}", + frequency_model="poisson", + frequency_lambda=0.05, + severity_model="lognormal", + severity_median="100 000", + severity_sigma=1.2, + severity_currency="USD", + tags="data-breach,baseline", + controls_json=( + '[{"id": "org:iam.mfa", "effectiveness_against_threat": 0.85}]' + if has_hsts else "[]" + ), + ) + + # Step 3: Validate + validation = ToolExecutor.call( + validate_crml_scenario, + context=context, + yaml_content=scenario_result["yaml"], + ) + + return AgentResult.ok( + context, + output={ + "domain": domain_name, + "scenarios_generated": 1, + "scenarios": [ + { + "name": scenario_result["scenario_name"], + "yaml": scenario_result["yaml"], + "validation": validation, + } + ], + "domain_intel": { + "domain": domain_name, + "has_hsts": has_hsts, + "headers_found": list(intel.get("http_headers", {}).keys()), + }, + }, + ) + + +# --------------------------------------------------------------------------- +# LLM-powered implementation +# --------------------------------------------------------------------------- + +class _LLMRiskScenarioAgent: + """ + Internal LLM-powered risk scenario generator using the ReAct loop. + + Not registered in AgentRegistry — accessed only through RiskScenarioAgent + when reasoning.mode == llm_react. + + The LLM follows this sequence: + 1. fetch_domain_intel -> understand the target company + 2. generate_crml_scenario -> create 3-5 tailored risk scenarios + 3. validate_crml_scenario -> verify each scenario is schema-valid + 4. STOP + summarize -> produce structured output + """ + + def execute(self, context: AgentContext) -> AgentResult: + from zak.core.llm.registry import get_llm_client + from zak.core.runtime.llm_agent import _build_openai_schema + from zak.core.tools.substrate import ToolExecutor + import json + + target_domain = context.metadata.get("target_domain", "unknown.com") + + available_tools = [fetch_domain_intel, generate_crml_scenario, validate_crml_scenario] + tools_schema = _build_openai_schema(available_tools) + + # LLM config from DSL + llm_cfg: dict[str, Any] = {} + if context.dsl.reasoning.llm: + llm_block = context.dsl.reasoning.llm + llm_cfg = ( + llm_block if isinstance(llm_block, dict) + else llm_block.model_dump(exclude_none=True) + ) + + client = get_llm_client( + provider=llm_cfg.get("provider"), + model=llm_cfg.get("model"), + ) + temperature = float(llm_cfg.get("temperature", 0.3)) + max_tokens = int(llm_cfg.get("max_tokens", 8192)) + max_iter = int(llm_cfg.get("max_iterations", 15)) + + system = f"""You are a cyber risk scenario generation agent for tenant '{context.tenant_id}'. + +TARGET DOMAIN: {target_domain} + +Your goal: Generate realistic, calibrated CRML (Cyber Risk Modeling Language) cyber risk scenarios +for the company behind the domain '{target_domain}'. + +CRML SCENARIO FORMAT (crml_scenario "1.0"): +Each scenario has: +- meta: name, description, tags, company_size, industries +- scenario.frequency: how often the threat occurs (poisson or hierarchical_gamma_poisson model) +- scenario.severity: conditional loss per event (lognormal, gamma, or mixture model) +- scenario.controls: optional relevant controls with effectiveness_against_threat (0.0-1.0) + +FREQUENCY MODEL GUIDANCE: +- Use 'poisson' with lambda parameter for FAIR-style models + - lambda = annual probability of occurrence (e.g. 0.05 = 5%/year, 0.15 = 15%/year) + - Ransomware: lambda 0.05-0.15, Phishing: 0.3-0.8, Data breach: 0.02-0.10 + - Larger companies face higher lambda for most threats +- Use 'hierarchical_gamma_poisson' with alpha_base/beta_base for QBER/Bayesian models + +SEVERITY MODEL GUIDANCE: +- Use 'lognormal' with median (human-readable, e.g. "250 000") and sigma (0.8-2.0) + - sigma < 1.0: low variability | sigma 1.0-1.5: moderate | sigma > 1.5: high variability + - Ransomware median: $200K-$2M depending on size | Data breach: $50K-$500K + - Larger companies: higher medians +- Currency should always be specified (default: USD) + +CONTROL ID FORMAT: +- Must be namespace:key (e.g. "org:iam.mfa", "org:net.firewall", "org:email.dmarc") +- Must NOT start with "attck:" (reserved) +- effectiveness_against_threat: 0.0-1.0 + +TOOL SEQUENCE: +1. Call fetch_domain_intel with the target domain to understand the company +2. Based on the intel, determine: industry, company size, likely tech stack, security posture +3. Generate 3-5 CRML scenarios covering different threat types: + - Ransomware attack + - Data breach / data exfiltration + - Phishing / social engineering (leading to credential compromise) + - Supply chain / third-party risk + - Business email compromise (BEC) + Choose scenarios most relevant to the identified industry. +4. For each scenario, call generate_crml_scenario with calibrated parameters +5. Validate each scenario with validate_crml_scenario +6. If validation fails, fix and regenerate + +OUTPUT FORMAT (your final response): +Return a JSON object with: +{{ + "domain": "{target_domain}", + "company_profile": {{ + "industry": "...", + "estimated_size": "smb|mid-market|enterprise|large-enterprise", + "security_posture_indicators": ["..."], + "risk_factors": ["..."] + }}, + "scenarios_generated": N, + "scenarios": [ + {{ + "name": "...", + "threat_type": "ransomware|data_breach|phishing|supply_chain|bec|...", + "model_style": "fair|qber", + "crml_yaml": "...", + "calibration_rationale": "Why these specific parameters were chosen" + }} + ], + "aggregate_risk_narrative": "Brief executive summary of the company's cyber risk profile" +}} + +Ground every parameter choice in the domain intelligence you gathered. Do not invent data. +Calibrate parameters to be realistic for the identified industry and company size.""" + + messages = [ + {"role": "system", "content": system}, + { + "role": "user", + "content": ( + f"Generate cyber risk scenarios for domain '{target_domain}'. " + f"Tenant: {context.tenant_id}. Environment: {context.environment}." + ), + }, + ] + + reasoning_trace: list[Any] = [] + total_usage: dict[str, Any] = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + def resolve(name: str) -> Any: + for fn in available_tools: + meta = getattr(fn, "_zak_tool", None) + if meta and meta.action_id == name: + return fn + return None + + for iteration in range(max_iter): + response = client.chat( + messages=messages, + tools=tools_schema, + temperature=temperature, + max_tokens=max_tokens, + ) + for k in total_usage: + total_usage[k] += response.usage.get(k, 0) + + if response.finish_reason == "stop" or not response.tool_calls: + conclusion = response.content or "Risk scenario generation complete." + reasoning_trace.append({ + "iteration": iteration + 1, "type": "conclusion", "content": conclusion, + }) + return AgentResult.ok( + context, + output={ + "summary": conclusion, + "reasoning_trace": reasoning_trace, + "iterations": iteration + 1, + "llm_usage": total_usage, + "provider": llm_cfg.get("provider", "openai"), + "model": llm_cfg.get("model"), + }, + ) + + tool_results = [] + for tc in response.tool_calls: + entry = { + "iteration": iteration + 1, "type": "tool_call", + "tool": tc.name, "arguments": tc.arguments, + } + reasoning_trace.append(entry) + fn = resolve(tc.name) + if fn is None: + err = {"error": f"Unknown tool: {tc.name}"} + entry["result"] = err + tool_results.append({ + "role": "tool", "tool_call_id": tc.id, + "content": json.dumps(err), + }) + continue + try: + result = ToolExecutor.call(fn, context=context, **tc.arguments) + entry["result"] = result + tool_results.append({ + "role": "tool", "tool_call_id": tc.id, + "content": json.dumps(result) if not isinstance(result, str) else result, + }) + except Exception as exc: + err = {"error": str(exc)} + entry["result"] = err + tool_results.append({ + "role": "tool", "tool_call_id": tc.id, + "content": json.dumps(err), + }) + + messages.append({ + "role": "assistant", + "content": response.content or "", + "tool_calls": [ # type: ignore[dict-item] + { + "id": tc.id, "type": "function", + "function": {"name": tc.name, "arguments": json.dumps(tc.arguments)}, + } + for tc in response.tool_calls + ], + }) + messages.extend(tool_results) + + return AgentResult.fail( + context, + errors=[f"LLM risk_scenario agent reached max_iterations ({max_iter}) without conclusion."], + ) diff --git a/zak/agents/risk_scenario/tools.py b/zak/agents/risk_scenario/tools.py new file mode 100644 index 0000000..1eea395 --- /dev/null +++ b/zak/agents/risk_scenario/tools.py @@ -0,0 +1,463 @@ +""" +Tools for the Risk Scenario Generator Agent. + +Provides domain intelligence gathering and CRML scenario generation/validation +capabilities that the LLM uses in its ReAct loop. +""" + +from __future__ import annotations + +import ipaddress +import json +import re +import socket +from typing import Any +from urllib.parse import urlparse + +import httpx +from bs4 import BeautifulSoup + +from zak.core.runtime.agent import AgentContext +from zak.core.tools.substrate import zak_tool + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_ALLOWED_SCHEMES = ("https", "http") + + +def _is_host_allowed(host: str) -> bool: + """Return True only if the host resolves to globally routable IPs (SSRF mitigation).""" + if not host or not host.strip(): + return False + host = host.strip().lower() + if host in ("localhost", "localhost.", "::1"): + return False + try: + infos = socket.getaddrinfo(host, None) + except (socket.gaierror, socket.herror, OSError): + return False + if not infos: + return False + for (_family, _type, _proto, _canon, sockaddr) in infos: + ip_str = sockaddr[0] if isinstance(sockaddr, (list, tuple)) else sockaddr + try: + ip = ipaddress.ip_address(ip_str) + if not ip.is_global: + return False + except ValueError: + return False + return True + + +def _validate_url(url: str) -> None: + """Validate URL scheme and host for safe external fetching.""" + parsed = urlparse(url) + scheme = (parsed.scheme or "").lower() + if scheme not in _ALLOWED_SCHEMES: + raise PermissionError( + f"Access denied: URL scheme must be one of {_ALLOWED_SCHEMES}, got {scheme or 'empty'}." + ) + netloc = parsed.netloc or parsed.path.split("/")[0] or "" + host = netloc.rsplit(":", 1)[0] if netloc else "" + if not host: + raise PermissionError("Access denied: URL has no host.") + if not _is_host_allowed(host): + raise PermissionError( + f"Access denied: '{host}' resolves to internal or disallowed addresses. " + "Only public website URLs are allowed." + ) + + +def _clean_html(html: str, max_chars: int = 15000) -> str: + """Extract and clean text from HTML content.""" + soup = BeautifulSoup(html, "html.parser") + for tag in soup(["script", "style", "nav", "footer", "header"]): + tag.decompose() + text = soup.get_text(separator="\n") + lines = (line.strip() for line in text.splitlines()) + chunks = (phrase.strip() for line in lines for phrase in line.split(" ")) + text = "\n".join(chunk for chunk in chunks if chunk) + return text[:max_chars] + + +# --------------------------------------------------------------------------- +# Tool: fetch_domain_intel +# --------------------------------------------------------------------------- + +@zak_tool( + name="fetch_domain_intel", + description=( + "Gather publicly available intelligence about a domain/company for risk scenario generation. " + "Fetches the company's website homepage and attempts to identify industry, size, tech stack, " + "and other risk-relevant context from the page content." + ), + action_id="fetch_domain_intel", + tags=["risk_scenario", "recon", "read"], +) +def fetch_domain_intel(context: AgentContext, domain: str) -> dict[str, Any]: + """ + Fetch publicly available information about a domain to inform risk scenario generation. + + Gathers the homepage content and extracts text for the LLM to analyse. The LLM + then uses this intelligence to calibrate realistic CRML scenario parameters. + + Args: + domain: The target domain (e.g. 'example.com'). Can also be a full URL. + + Returns: + Dict with domain, url, page_text, and any extracted metadata. + """ + # Normalise: accept both "example.com" and "https://example.com" + if not domain.startswith(("http://", "https://")): + url = f"https://{domain}" + else: + url = domain + + parsed = urlparse(url) + clean_domain = parsed.netloc or parsed.hostname or domain + + _validate_url(url) + + result: dict[str, Any] = { + "domain": clean_domain, + "url": url, + "page_text": "", + "http_headers": {}, + "error": None, + } + + try: + resp = httpx.get( + url, + timeout=15.0, + follow_redirects=True, + headers={"User-Agent": "ZAK-RiskScenarioAgent/1.0"}, + ) + resp.raise_for_status() + + # Capture useful headers for tech-stack inference + interesting_headers = [ + "server", "x-powered-by", "x-frame-options", + "content-security-policy", "strict-transport-security", + "x-content-type-options", "x-xss-protection", + ] + result["http_headers"] = { + k: v for k, v in resp.headers.items() + if k.lower() in interesting_headers + } + + result["page_text"] = _clean_html(resp.text) + + except PermissionError: + raise + except Exception as exc: + result["error"] = f"Failed to fetch {url}: {str(exc)}" + + return result + + +# --------------------------------------------------------------------------- +# Tool: generate_crml_scenario +# --------------------------------------------------------------------------- + +@zak_tool( + name="generate_crml_scenario", + description=( + "Generate a valid CRML (Cyber Risk Modeling Language) scenario YAML document. " + "Accepts structured risk parameters (name, description, frequency model, severity model, " + "controls, metadata) and produces a schema-valid crml_scenario YAML string." + ), + action_id="generate_crml_scenario", + tags=["risk_scenario", "crml", "write"], +) +def generate_crml_scenario( + context: AgentContext, + name: str, + description: str, + frequency_model: str, + frequency_lambda: float, + severity_model: str, + severity_median: str, + severity_sigma: float, + severity_currency: str = "USD", + tags: str = "", + company_size: str = "", + industries: str = "", + controls_json: str = "", + frequency_basis: str = "per_organization_per_year", + author: str = "ZAK Risk Scenario Agent", + frequency_alpha_base: float = 0.0, + frequency_beta_base: float = 0.0, + severity_components_json: str = "", +) -> dict[str, Any]: + """ + Build a CRML scenario YAML document from structured parameters. + + The LLM calls this tool after analysing domain intelligence to produce + a properly formatted CRML scenario. Supports FAIR-style (poisson + lognormal) + and QBER-style (hierarchical_gamma_poisson + mixture) models. + + Args: + name: Scenario name (kebab-case recommended). + description: Human-readable description of the risk scenario. + frequency_model: One of: poisson, hierarchical_gamma_poisson, gamma. + frequency_lambda: Poisson lambda (events/year). Used for poisson model. + severity_model: One of: lognormal, gamma, mixture. + severity_median: Median loss amount as string (e.g. "250 000"). + severity_sigma: Lognormal sigma (loss variability). Typical: 0.8-2.0. + severity_currency: ISO 4217 currency code (default: USD). + tags: Comma-separated tags (e.g. "ransomware,encryption"). + company_size: Comma-separated sizes (e.g. "smb,enterprise"). + industries: Comma-separated industries (e.g. "healthcare,finance"). + controls_json: JSON array of controls, each with id and effectiveness_against_threat. + frequency_basis: per_organization_per_year or per_asset_unit_per_year. + author: Scenario author. + frequency_alpha_base: Gamma shape for hierarchical_gamma_poisson model. + frequency_beta_base: Gamma scale for hierarchical_gamma_poisson model. + severity_components_json: JSON array of mixture components for mixture severity model. + + Returns: + Dict with 'yaml' (the CRML YAML string) and 'scenario_name'. + """ + import yaml # type: ignore[import-untyped] + + # --- Build meta --- + meta: dict[str, Any] = { + "name": name, + "version": "1.0", + "description": description, + "author": author, + } + if tags: + meta["tags"] = [t.strip() for t in tags.split(",") if t.strip()] + if company_size: + meta["company_size"] = [s.strip() for s in company_size.split(",") if s.strip()] + if industries: + meta["industries"] = [i.strip() for i in industries.split(",") if i.strip()] + + # --- Build frequency --- + frequency: dict[str, Any] = { + "basis": frequency_basis, + "model": frequency_model, + } + if frequency_model == "poisson": + frequency["parameters"] = {"lambda": frequency_lambda} + elif frequency_model == "hierarchical_gamma_poisson": + frequency["parameters"] = { + "alpha_base": frequency_alpha_base or 1.5, + "beta_base": frequency_beta_base or 1.5, + } + elif frequency_model == "gamma": + frequency["parameters"] = { + "shape": frequency_alpha_base or 2.0, + "scale": frequency_beta_base or 1.0, + } + + # --- Build severity --- + severity: dict[str, Any] = {"model": severity_model} + + if severity_model == "lognormal": + severity["parameters"] = { + "median": severity_median, + "currency": severity_currency, + "sigma": severity_sigma, + } + elif severity_model == "gamma": + severity["parameters"] = { + "shape": severity_sigma, + "scale": severity_median, + "currency": severity_currency, + } + elif severity_model == "mixture": + severity["parameters"] = {} + if severity_components_json: + try: + severity["components"] = json.loads(severity_components_json) + except json.JSONDecodeError: + severity["components"] = [ + {"lognormal": {"weight": 0.7, "median": severity_median, "currency": severity_currency, "sigma": severity_sigma}}, + {"gamma": {"weight": 0.3, "shape": 2.5, "scale": "10 000", "currency": severity_currency}}, + ] + else: + severity["components"] = [ + {"lognormal": {"weight": 0.7, "median": severity_median, "currency": severity_currency, "sigma": severity_sigma}}, + {"gamma": {"weight": 0.3, "shape": 2.5, "scale": "10 000", "currency": severity_currency}}, + ] + + # --- Build controls --- + controls = [] + if controls_json: + try: + controls = json.loads(controls_json) + except json.JSONDecodeError: + pass + + # --- Assemble scenario --- + scenario_payload: dict[str, Any] = { + "frequency": frequency, + "severity": severity, + } + if controls: + scenario_payload["controls"] = controls + + doc = { + "crml_scenario": "1.0", + "meta": meta, + "scenario": scenario_payload, + } + + yaml_str = yaml.dump(doc, default_flow_style=False, sort_keys=False, allow_unicode=True) + + return { + "yaml": yaml_str, + "scenario_name": name, + } + + +# --------------------------------------------------------------------------- +# Tool: validate_crml_scenario +# --------------------------------------------------------------------------- + +@zak_tool( + name="validate_crml_scenario", + description=( + "Validate a CRML scenario YAML string against the CRML specification. " + "Checks required fields, frequency/severity model validity, control ID formats, " + "and parameter ranges. Returns validation result with any errors found." + ), + action_id="validate_crml_scenario", + tags=["risk_scenario", "crml", "validation"], +) +def validate_crml_scenario(context: AgentContext, yaml_content: str) -> dict[str, Any]: + """ + Validate a CRML scenario YAML string for structural and semantic correctness. + + Checks: + - Required top-level fields (crml_scenario, meta.name, scenario) + - Frequency model is supported and has required parameters + - Severity model is supported and has required parameters + - Control IDs follow the namespace:key format + - Parameter value ranges are sensible + + Args: + yaml_content: The CRML scenario as a YAML string. + + Returns: + Dict with 'valid' (bool), 'errors' (list[str]), and 'warnings' (list[str]). + """ + import yaml + + errors: list[str] = [] + warnings: list[str] = [] + + # --- Parse YAML --- + try: + doc = yaml.safe_load(yaml_content) + except yaml.YAMLError as exc: + return {"valid": False, "errors": [f"YAML parse error: {exc}"], "warnings": []} + + if not isinstance(doc, dict): + return {"valid": False, "errors": ["Document root must be a mapping"], "warnings": []} + + # --- Top-level fields --- + if "crml_scenario" not in doc: + errors.append("Missing required field: crml_scenario (must be '1.0')") + elif str(doc["crml_scenario"]) != "1.0": + warnings.append(f"crml_scenario version '{doc['crml_scenario']}' — only '1.0' is currently supported") + + meta = doc.get("meta", {}) + if not meta.get("name"): + errors.append("Missing required field: meta.name") + + scenario = doc.get("scenario") + if not scenario: + errors.append("Missing required field: scenario") + return {"valid": len(errors) == 0, "errors": errors, "warnings": warnings} + + # --- Frequency validation --- + freq = scenario.get("frequency") + if not freq: + errors.append("Missing required field: scenario.frequency") + else: + valid_freq_models = {"poisson", "hierarchical_gamma_poisson", "gamma"} + model = freq.get("model", "") + if model not in valid_freq_models: + errors.append(f"Invalid frequency model '{model}'. Must be one of: {valid_freq_models}") + + basis = freq.get("basis", "per_organization_per_year") + valid_bases = {"per_organization_per_year", "per_asset_unit_per_year"} + if basis not in valid_bases: + errors.append(f"Invalid frequency basis '{basis}'. Must be one of: {valid_bases}") + + params = freq.get("parameters", {}) + if model == "poisson": + lam = params.get("lambda") + if lam is None: + errors.append("Poisson frequency model requires 'lambda' parameter") + elif isinstance(lam, (int, float)) and lam < 0: + errors.append(f"Poisson lambda must be non-negative, got {lam}") + elif isinstance(lam, (int, float)) and lam > 100: + warnings.append(f"Poisson lambda={lam} is unusually high — verify this is intentional") + elif model == "hierarchical_gamma_poisson": + if "alpha_base" not in params: + errors.append("hierarchical_gamma_poisson requires 'alpha_base' parameter") + if "beta_base" not in params: + errors.append("hierarchical_gamma_poisson requires 'beta_base' parameter") + + # --- Severity validation --- + sev = scenario.get("severity") + if not sev: + errors.append("Missing required field: scenario.severity") + else: + valid_sev_models = {"lognormal", "gamma", "mixture"} + model = sev.get("model", "") + if model not in valid_sev_models: + errors.append(f"Invalid severity model '{model}'. Must be one of: {valid_sev_models}") + + params = sev.get("parameters", {}) + if model == "lognormal": + has_median = "median" in params + has_mu = "mu" in params + if not has_median and not has_mu: + errors.append("Lognormal severity requires 'median' or 'mu' parameter") + if "sigma" not in params: + errors.append("Lognormal severity requires 'sigma' parameter") + sigma = params.get("sigma") + if isinstance(sigma, (int, float)) and sigma <= 0: + errors.append(f"Lognormal sigma must be positive, got {sigma}") + elif model == "mixture": + components = sev.get("components", []) + if not components: + errors.append("Mixture severity model requires at least one component") + else: + total_weight = 0.0 + for comp in components: + if isinstance(comp, dict): + for dist_type, dist_params in comp.items(): + w = dist_params.get("weight", 0) + total_weight += float(w) + if abs(total_weight - 1.0) > 0.01: + warnings.append(f"Mixture component weights sum to {total_weight}, expected ~1.0") + + # --- Controls validation --- + controls = scenario.get("controls", []) + control_id_pattern = re.compile(r"^[a-z][a-z0-9_-]{0,31}:[^\s]{1,223}$") + for ctrl in controls: + if isinstance(ctrl, dict): + ctrl_id = ctrl.get("id", "") + if ctrl_id.startswith("attck:"): + errors.append(f"Control ID '{ctrl_id}' must not start with 'attck:' (reserved for attacks)") + elif ctrl_id and not control_id_pattern.match(ctrl_id): + warnings.append(f"Control ID '{ctrl_id}' does not match recommended format namespace:key") + eff = ctrl.get("effectiveness_against_threat") + if eff is not None and (float(eff) < 0.0 or float(eff) > 1.0): + errors.append(f"Control effectiveness must be 0.0-1.0, got {eff} for '{ctrl_id}'") + + return { + "valid": len(errors) == 0, + "errors": errors, + "warnings": warnings, + } diff --git a/zak/agents/vuln_triage/agent.py b/zak/agents/vuln_triage/agent.py index 7b8af72..dc210f6 100644 --- a/zak/agents/vuln_triage/agent.py +++ b/zak/agents/vuln_triage/agent.py @@ -18,6 +18,8 @@ from __future__ import annotations +from typing import Any + from zak.core.dsl.schema import ReasoningMode from zak.core.runtime.agent import AgentContext, AgentResult, BaseAgent from zak.core.runtime.llm_agent import LLMAgent @@ -58,16 +60,16 @@ def execute(self, context: AgentContext) -> AgentResult: def _execute_deterministic(self, context: AgentContext) -> AgentResult: tenant_id = context.tenant_id vulns = ( - self._adapter.get_nodes(tenant_id=tenant_id, node_type="vulnerability") # type: ignore[union-attr] + self._adapter.get_nodes(tenant_id=tenant_id, node_type="vulnerability") # type: ignore[attr-defined] if self._adapter is not None else [] ) assets = ( - self._adapter.get_nodes(tenant_id=tenant_id, node_type="asset") # type: ignore[union-attr] + self._adapter.get_nodes(tenant_id=tenant_id, node_type="asset") # type: ignore[attr-defined] if self._adapter is not None else [] ) asset_crit = {a["node_id"]: a.get("criticality", "medium") for a in assets} - triaged: list[dict] = [] + triaged: list[dict[str, Any]] = [] for vuln in vulns: severity = vuln.get("severity", "medium").lower() @@ -119,7 +121,7 @@ class _LLMVulnTriageAgent(LLMAgent): """ @property - def tools(self) -> list: + def tools(self) -> list[Any]: from zak.core.tools.builtins import ( list_vulnerabilities, list_assets, diff --git a/zak/cli/main.py b/zak/cli/main.py index 6bd99d9..a43c1a5 100644 --- a/zak/cli/main.py +++ b/zak/cli/main.py @@ -25,6 +25,7 @@ from zak.core.edition import Edition, EditionError, get_edition from zak.agents import load_all_agents as _load_all_agents +from zak.sif.schema.nodes import Criticality, Environment, ExposureLevel, Severity console = Console() @@ -71,20 +72,26 @@ def quickstart(with_llm: bool) -> None: AssetNode( node_id="web-server-prod", asset_type="server", - criticality="critical", - environment="production", - exposure_level="internet_facing", + criticality=Criticality.CRITICAL, + environment=Environment.PRODUCTION, + exposure_level=ExposureLevel.INTERNET_FACING, owner="platform-team", source="quickstart-demo", + valid_to=None, + confidence=1.0, + risk_score=0.0, ), AssetNode( node_id="internal-api", asset_type="application", - criticality="high", - environment="production", - exposure_level="internal", + criticality=Criticality.HIGH, + environment=Environment.PRODUCTION, + exposure_level=ExposureLevel.INTERNAL, owner="backend-team", source="quickstart-demo", + valid_to=None, + confidence=1.0, + risk_score=0.0, ), ] sample_vulns = [ @@ -92,17 +99,22 @@ def quickstart(with_llm: bool) -> None: node_id="CVE-2024-1234", vuln_type="cve", cve_id="CVE-2024-1234", - severity="critical", + severity=Severity.CRITICAL, exploitability=0.9, cvss_score=9.8, source="quickstart-demo", + valid_to=None, + confidence=1.0, ), VulnerabilityNode( node_id="MISCONFIG-001", vuln_type="misconfiguration", - severity="medium", + severity=Severity.MEDIUM, exploitability=0.4, + cvss_score=None, source="quickstart-demo", + valid_to=None, + confidence=1.0, ), ] sample_controls = [ @@ -112,6 +124,8 @@ def quickstart(with_llm: bool) -> None: effectiveness=0.7, automated=True, source="quickstart-demo", + valid_to=None, + confidence=1.0, ), ] @@ -478,7 +492,7 @@ def run(path: str, tenant: str, env: str, meta: list[str]) -> None: # Inject graph adapter for agents that need it (check constructor signature) import inspect - sig = inspect.signature(agent_cls.__init__) + sig = inspect.signature(agent_cls.__init__) # type: ignore[misc] if "adapter" in sig.parameters: from zak.sif.graph.factory import create_adapter adapter = create_adapter() diff --git a/zak/cli/templates.py b/zak/cli/templates.py index 1790db4..34cd93a 100644 --- a/zak/cli/templates.py +++ b/zak/cli/templates.py @@ -402,4 +402,119 @@ def execute(self, context: AgentContext) -> AgentResult: ], default_denied_actions=["delete_branch", "force_push"], ), + "risk_scenario": DomainTemplate( + yaml_template="""\ +agent: + id: {agent_id} + name: "{agent_name}" + domain: risk_scenario + version: "1.0.0" + +intent: + goal: "Generate calibrated CRML cyber risk scenarios for a target domain" + success_criteria: + - "At least 3 CRML scenarios generated" + - "All scenarios pass CRML schema validation" + - "Parameters calibrated to industry and company size" + priority: high + +reasoning: + mode: llm_react + autonomy_level: bounded + confidence_threshold: 0.75 + llm: + provider: openai + model: gpt-4o + temperature: 0.3 + max_iterations: 15 + max_tokens: 8192 + +capabilities: + tools: + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + data_access: + - web_public + +boundaries: + risk_budget: low + allowed_actions: + - agent_execute + - fetch_domain_intel + - generate_crml_scenario + - validate_crml_scenario + denied_actions: + - write_risk_node + - delete_asset + environment_scope: + - production + - staging + +safety: + guardrails: + - no_destructive_actions + - require_confidence_threshold + sandbox_profile: standard + audit_level: standard +""", + python_template="""\ +\"\"\" +{agent_name} — ZAK Risk Scenario Generator Agent. + +Auto-generated by `zak init`. Generates CRML cyber risk scenarios for a target domain. +\"\"\" + +from __future__ import annotations + +from zak.core.runtime.agent import AgentContext, AgentResult, BaseAgent +from zak.core.runtime.registry import register_agent +from zak.core.tools.substrate import ToolExecutor +from zak.agents.risk_scenario import scenario_tools as tools # noqa: F401 + + +@register_agent( + domain="risk_scenario", + description="{agent_name}", + version="1.0.0", +) +class {class_name}(BaseAgent): + \"\"\" + {agent_name} + + Generates CRML cyber risk scenarios from domain intelligence. + \"\"\" + + def execute(self, context: AgentContext) -> AgentResult: + target_domain = context.metadata.get("target_domain", "") + + # 1. Gather domain intelligence + intel = ToolExecutor.call( + tools.fetch_domain_intel, context=context, domain=target_domain, + ) + + # 2. Generate a baseline scenario + scenario = ToolExecutor.call( + tools.generate_crml_scenario, context=context, + name=f"{{target_domain}}-data-breach", + description=f"Data breach scenario for {{target_domain}}", + frequency_model="poisson", frequency_lambda=0.05, + severity_model="lognormal", severity_median="100 000", severity_sigma=1.2, + ) + + # 3. Validate the scenario + validation = ToolExecutor.call( + tools.validate_crml_scenario, context=context, + yaml_content=scenario["yaml"], + ) + + return AgentResult.ok(context, output={{ + "scenarios": [scenario], + "validation": validation, + }}) +""", + default_tools=["fetch_domain_intel", "generate_crml_scenario", "validate_crml_scenario"], + default_allowed_actions=["fetch_domain_intel", "generate_crml_scenario", "validate_crml_scenario"], + default_denied_actions=["write_risk_node", "delete_asset"], + ), } diff --git a/zak/core/dsl/parser.py b/zak/core/dsl/parser.py index d405fb8..dfd272b 100644 --- a/zak/core/dsl/parser.py +++ b/zak/core/dsl/parser.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import Any -import yaml +import yaml # type: ignore[import-untyped] from pydantic import ValidationError from zak.core.dsl.schema import AgentDSL diff --git a/zak/core/dsl/schema.py b/zak/core/dsl/schema.py index 85aa57d..1d2e268 100644 --- a/zak/core/dsl/schema.py +++ b/zak/core/dsl/schema.py @@ -24,6 +24,7 @@ class Domain(str, Enum): AI_SECURITY = "ai_security" RISK_QUANT = "risk_quant" SUPPLY_CHAIN = "supply_chain" + RISK_SCENARIO = "risk_scenario" COMPLIANCE = "compliance" # Enterprise domains API_SECURITY = "api_security" @@ -260,9 +261,9 @@ class AgentDSL(BaseModel): agent: AgentIdentity intent: AgentIntent reasoning: ReasoningConfig - capabilities: CapabilitiesConfig = Field(default_factory=CapabilitiesConfig) - boundaries: BoundariesConfig = Field(default_factory=BoundariesConfig) - safety: SafetyConfig = Field(default_factory=SafetyConfig) + capabilities: CapabilitiesConfig = Field(default_factory=lambda: CapabilitiesConfig(tools=[], data_access=[], graph_access=[])) + boundaries: BoundariesConfig = Field(default_factory=lambda: BoundariesConfig(risk_budget=RiskBudget.MEDIUM, allowed_actions=[], denied_actions=[], environment_scope=[], approval_gates=[])) + safety: SafetyConfig = Field(default_factory=lambda: SafetyConfig(guardrails=[], sandbox_profile=SandboxProfile.STANDARD, audit_level=AuditLevel.STANDARD)) @model_validator(mode="after") def offensive_agents_require_isolated_sandbox(self) -> AgentDSL: @@ -285,7 +286,13 @@ def llm_react_requires_llm_config(self) -> AgentDSL: if self.reasoning.llm is None: # Auto-populate with defaults so the field is never missing self.reasoning = self.reasoning.model_copy( - update={"llm": LLMConfig()} + update={"llm": LLMConfig( + provider="openai", + model="gpt-4o", + temperature=0.2, + max_iterations=10, + max_tokens=4096, + )} ) return self diff --git a/zak/core/llm/anthropic_client.py b/zak/core/llm/anthropic_client.py index d1bc187..3e281dc 100644 --- a/zak/core/llm/anthropic_client.py +++ b/zak/core/llm/anthropic_client.py @@ -45,10 +45,10 @@ def chat( import httpx if self._client is None: http_client = httpx.Client(verify=False) - kwargs: dict[str, Any] = {"api_key": self.api_key, "http_client": http_client} + init_kwargs: dict[str, Any] = {"api_key": self.api_key, "http_client": http_client} if self.base_url: - kwargs["base_url"] = self.base_url - self._client = anthropic.Anthropic(**kwargs) + init_kwargs["base_url"] = self.base_url + self._client = anthropic.Anthropic(**init_kwargs) client = self._client # Separate system message from conversation history diff --git a/zak/core/llm/local.py b/zak/core/llm/local.py index 538415b..56640ad 100644 --- a/zak/core/llm/local.py +++ b/zak/core/llm/local.py @@ -35,10 +35,11 @@ def __init__( base_url: str | None = None, ) -> None: self.model = model or os.getenv("LLM_MODEL", "llama3.1:8b") - self.base_url = ( + base_url_str = ( base_url or os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") - ).rstrip("/") + ) + self.base_url = base_url_str.rstrip("/") if base_url_str else "" def chat( self, diff --git a/zak/core/llm/openai_client.py b/zak/core/llm/openai_client.py index 4f55429..5285f96 100644 --- a/zak/core/llm/openai_client.py +++ b/zak/core/llm/openai_client.py @@ -45,10 +45,10 @@ def chat( ) from exc if self._client is None: - kwargs: dict[str, Any] = {"api_key": self.api_key} + init_kwargs: dict[str, Any] = {"api_key": self.api_key} if self.base_url: - kwargs["base_url"] = self.base_url - self._client = openai.OpenAI(**kwargs) + init_kwargs["base_url"] = self.base_url + self._client = openai.OpenAI(**init_kwargs) client = self._client kwargs: dict[str, Any] = dict( diff --git a/zak/core/runtime/llm_agent.py b/zak/core/runtime/llm_agent.py index c995127..03c1fcc 100644 --- a/zak/core/runtime/llm_agent.py +++ b/zak/core/runtime/llm_agent.py @@ -249,9 +249,9 @@ def execute(self, context: AgentContext) -> AgentResult: tool_fn = self._resolve_tool(tool_call.name) if tool_fn is None: available = [ - getattr(t, "_zak_tool", None).action_id + meta.action_id for t in self.tools - if getattr(t, "_zak_tool", None) + if (meta := getattr(t, "_zak_tool", None)) ] err = { "error": ( @@ -514,9 +514,9 @@ def execute_stream( tool_fn = self._resolve_tool(tool_call.name) if tool_fn is None: available = [ - getattr(t, "_zak_tool", None).action_id + meta.action_id for t in self.tools - if getattr(t, "_zak_tool", None) + if (meta := getattr(t, "_zak_tool", None)) ] err = { "error": ( diff --git a/zak/core/runtime/registry.py b/zak/core/runtime/registry.py index a9d3814..99f0b43 100644 --- a/zak/core/runtime/registry.py +++ b/zak/core/runtime/registry.py @@ -251,9 +251,9 @@ def decorator(cls: type) -> type: override=override, ) # Attach metadata to the class itself for introspection - cls._zak_domain = domain - cls._zak_version = version - cls._zak_edition = edition + cls._zak_domain = domain # type: ignore[attr-defined] + cls._zak_version = version # type: ignore[attr-defined] + cls._zak_edition = edition # type: ignore[attr-defined] return cls return decorator diff --git a/zak/core/tools/builtins.py b/zak/core/tools/builtins.py index 56d627c..4b11619 100644 --- a/zak/core/tools/builtins.py +++ b/zak/core/tools/builtins.py @@ -42,7 +42,7 @@ def _get_adapter() -> Any: def read_asset(context: AgentContext, asset_id: str) -> Optional[dict[str, Any]]: """Read a single asset node from the SIF graph.""" adapter = _get_adapter() - return adapter.get_node(tenant_id=context.tenant_id, node_type="asset", node_id=asset_id) + return adapter.get_node(tenant_id=context.tenant_id, node_type="asset", node_id=asset_id) # type: ignore[no-any-return] @zak_tool( @@ -54,7 +54,7 @@ def read_asset(context: AgentContext, asset_id: str) -> Optional[dict[str, Any]] def list_assets(context: AgentContext) -> list[dict[str, Any]]: """List all assets for the current tenant.""" adapter = _get_adapter() - return adapter.get_nodes(tenant_id=context.tenant_id, node_type="asset") + return adapter.get_nodes(tenant_id=context.tenant_id, node_type="asset") # type: ignore[no-any-return] @zak_tool( @@ -66,7 +66,7 @@ def list_assets(context: AgentContext) -> list[dict[str, Any]]: def list_vulnerabilities(context: AgentContext) -> list[dict[str, Any]]: """List all vulnerability nodes for the current tenant.""" adapter = _get_adapter() - return adapter.get_nodes(tenant_id=context.tenant_id, node_type="vulnerability") + return adapter.get_nodes(tenant_id=context.tenant_id, node_type="vulnerability") # type: ignore[no-any-return] @zak_tool( @@ -78,7 +78,7 @@ def list_vulnerabilities(context: AgentContext) -> list[dict[str, Any]]: def list_vendors(context: AgentContext) -> list[dict[str, Any]]: """List all vendor nodes for the current tenant.""" adapter = _get_adapter() - return adapter.get_nodes(tenant_id=context.tenant_id, node_type="vendor") + return adapter.get_nodes(tenant_id=context.tenant_id, node_type="vendor") # type: ignore[no-any-return] @zak_tool( @@ -90,7 +90,7 @@ def list_vendors(context: AgentContext) -> list[dict[str, Any]]: def list_controls(context: AgentContext) -> list[dict[str, Any]]: """List all security control nodes for the current tenant.""" adapter = _get_adapter() - return adapter.get_nodes(tenant_id=context.tenant_id, node_type="control") + return adapter.get_nodes(tenant_id=context.tenant_id, node_type="control") # type: ignore[no-any-return] @zak_tool( @@ -102,7 +102,7 @@ def list_controls(context: AgentContext) -> list[dict[str, Any]]: def list_identities(context: AgentContext) -> list[dict[str, Any]]: """List all identity nodes for the current tenant.""" adapter = _get_adapter() - return adapter.get_nodes(tenant_id=context.tenant_id, node_type="identity") + return adapter.get_nodes(tenant_id=context.tenant_id, node_type="identity") # type: ignore[no-any-return] @zak_tool( @@ -114,7 +114,7 @@ def list_identities(context: AgentContext) -> list[dict[str, Any]]: def list_risks(context: AgentContext) -> list[dict[str, Any]]: """List all risk nodes for the current tenant.""" adapter = _get_adapter() - return adapter.get_nodes(tenant_id=context.tenant_id, node_type="risk") + return adapter.get_nodes(tenant_id=context.tenant_id, node_type="risk") # type: ignore[no-any-return] @zak_tool( @@ -126,7 +126,7 @@ def list_risks(context: AgentContext) -> list[dict[str, Any]]: def list_ai_models(context: AgentContext) -> list[dict[str, Any]]: """List all AI model nodes for the current tenant.""" adapter = _get_adapter() - return adapter.get_nodes(tenant_id=context.tenant_id, node_type="ai_model") + return adapter.get_nodes(tenant_id=context.tenant_id, node_type="ai_model") # type: ignore[no-any-return] # --------------------------------------------------------------------------- diff --git a/zak/core/tools/orchestration.py b/zak/core/tools/orchestration.py index bfb7129..d89e1eb 100644 --- a/zak/core/tools/orchestration.py +++ b/zak/core/tools/orchestration.py @@ -93,7 +93,7 @@ def spawn_agent( except Exception as exc: return {"domain": domain, "error": str(exc)} - sig = inspect.signature(agent_cls.__init__) + sig = inspect.signature(agent_cls.__init__) # type: ignore[misc] if "adapter" in sig.parameters: try: from zak.sif.graph.adapter import KuzuAdapter diff --git a/zak/core/tools/substrate.py b/zak/core/tools/substrate.py index 2d1ae58..270e3f1 100644 --- a/zak/core/tools/substrate.py +++ b/zak/core/tools/substrate.py @@ -59,7 +59,7 @@ class ToolRegistry: _lock: threading.Lock = threading.Lock() def __init__(self) -> None: - self._tools: dict[str, tuple[ToolMetadata, Callable]] = {} + self._tools: dict[str, tuple[ToolMetadata, Callable[..., Any]]] = {} @classmethod def get(cls) -> ToolRegistry: @@ -69,10 +69,10 @@ def get(cls) -> ToolRegistry: cls._instance = cls() return cls._instance - def register(self, metadata: ToolMetadata, fn: Callable) -> None: + def register(self, metadata: ToolMetadata, fn: Callable[..., Any]) -> None: self._tools[metadata.action_id] = (metadata, fn) - def get_tool(self, action_id: str) -> tuple[ToolMetadata, Callable] | None: + def get_tool(self, action_id: str) -> tuple[ToolMetadata, Callable[..., Any]] | None: return self._tools.get(action_id) def all_tools(self) -> list[ToolMetadata]: @@ -99,7 +99,7 @@ def zak_tool( description: str = "", action_id: str | None = None, tags: list[str] | None = None, -) -> Callable: +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: """ Decorator that registers a function as a ZAK tool. @@ -120,7 +120,7 @@ def read_asset(context: AgentContext, asset_id: str) -> dict: """ resolved_action_id = action_id or name.lower().replace(" ", "_") - def decorator(fn: Callable) -> Callable: + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: meta = ToolMetadata( name=name, description=description or (inspect.getdoc(fn) or "").split("\n")[0], @@ -158,7 +158,7 @@ class ToolExecutor: @classmethod def call( cls, - tool_fn: Callable, + tool_fn: Callable[..., Any], context: AgentContext, **kwargs: Any, ) -> Any: diff --git a/zak/sif/graph/adapter.py b/zak/sif/graph/adapter.py index 0528e6a..9bed94d 100644 --- a/zak/sif/graph/adapter.py +++ b/zak/sif/graph/adapter.py @@ -30,9 +30,9 @@ from typing import Any, Optional try: - from neo4j import GraphDatabase as _Neo4jGD # type: ignore[import-untyped] + from neo4j import GraphDatabase as _Neo4jGD except ImportError: - _Neo4jGD = None # type: ignore[assignment] + _Neo4jGD = None from zak.sif.schema.nodes import ( AIModelNode, diff --git a/zak/sif/telemetry/ingestor.py b/zak/sif/telemetry/ingestor.py index 1c8b352..09cea54 100644 --- a/zak/sif/telemetry/ingestor.py +++ b/zak/sif/telemetry/ingestor.py @@ -76,7 +76,10 @@ def _handle_asset_discovered( environment=event.get("environment", "production"), owner=event.get("owner"), exposure_level=event.get("exposure_level", "internal"), + risk_score=_safe_float(event.get("risk_score", 0.0), 0.0), source=event.get("source", "telemetry"), + valid_to=None, + confidence=0.95, ) self._adapter.upsert_node(tenant_id, node) @@ -91,6 +94,8 @@ def _handle_vulnerability_found( exploitability=_safe_float(event.get("exploitability", 0.5), 0.5), cvss_score=event.get("cvss_score"), source=event.get("source", "telemetry"), + valid_to=None, + confidence=0.95, ) self._adapter.upsert_node(tenant_id, vuln) @@ -115,6 +120,8 @@ def _handle_control_updated( effectiveness=_safe_float(event.get("effectiveness", 0.5), 0.5), automated=event.get("automated", True), source=event.get("source", "telemetry"), + valid_to=None, + confidence=0.95, ) self._adapter.upsert_node(tenant_id, node) @@ -128,5 +135,7 @@ def _handle_vendor_assessed( risk_score=_safe_float(event.get("risk_score", 0.0), 0.0), last_assessed=datetime.now(timezone.utc), source=event.get("source", "telemetry"), + valid_to=None, + confidence=0.95, ) self._adapter.upsert_node(tenant_id, node)