diff --git a/README.md b/README.md index b5d3f46..c916eee 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ In modern LLM-powered systems, agentic workflows are becoming increasingly compl |---------|-------------| | **Workflow Extraction** | Extract agent-based workflows from code using AST-based static parsing | | **Threat Vector Visualization** | View agent-to-agent, agent-to-tool, and chained connections in an interactive graph | -| **Automated Threat Assessment** | Generate comprehensive threat assessment report using established threat modeling frameworks for AI agents like MAESTRO| +| **Automated Threat Assessment** | Generate comprehensive threat assessment report using established threat modeling frameworks for AI agents like MAESTRO or STRIDE | | **Framework Agnostic** | Works with all major LLM orchestration frameworks | | **Developer Friendly** | Simple CLI, extensible SDK, and clean JSON exports | @@ -78,7 +78,7 @@ Each framework has its own AST-based static parser to extract: ## Security Analysis -Agent Wiz currently supports [**MAESTRO**](https://cloudsecurityalliance.org/blog/2025/02/06/agentic-ai-threat-modeling-framework-maestro) as its primary threat modeling framework. It evaluates agent workflows against the following structure: +Agent Wiz currently supports [**MAESTRO**](https://cloudsecurityalliance.org/blog/2025/02/06/agentic-ai-threat-modeling-framework-maestro) and [**STRIDE**](https://www.practical-devsecops.com/what-is-stride-threat-model/) threat modeling methodologies. MAESTRO remains the default; pass `--methodology stride` to switch the analyzer to STRIDE. MAESTRO evaluates agent workflows against the following structure: - **M**ission: Defining the system purpose and security objectives - **A**ssets: Inventorying critical components (agents, tools, data flows) @@ -88,6 +88,8 @@ Agent Wiz currently supports [**MAESTRO**](https://cloudsecurityalliance.org/blo - **R**isks: Calculating impact and likelihood of security events - **O**perations: Assessing runtime security considerations +STRIDE reports organize findings across the six foundational categories: Spoofing, Tampering, Repudiation, Information Disclosure, Denial of Service, and Elevation of Privilege—highlighting how each manifests within the extracted agentic workflow and the mitigations to prioritize. + Sample threat modelling report generated: Threat Modeling Report @@ -96,7 +98,7 @@ Sample threat modelling report generated: You can also add this line to your `.bashrc`, `.zshrc`, or environment setup script for persistent use. -🧪 More threat models analysis (STRIDE, PASTA, LINDDUN, etc.) are under development. +🧪 Additional threat modeling frameworks (PASTA, LINDDUN, etc.) are under development. ## Installation @@ -142,7 +144,7 @@ This will generate an html d3 based visualisation of the agentic workflow. The ` ### 3. Analyze against Threat Modeling ```bash -agent-wiz analyze --input agentchat_graph.json +agent-wiz analyze --input agentchat_graph.json --methodology (Default to maestro can be changed to stride) ``` This will generate a report like: `autogen_report.md` based on the provided graph and threat modeling frameworks. @@ -169,7 +171,8 @@ Planned features (Not in any paricular order) - [x] Generate standardized JSON graph representations of agent flows - [x] CLI interfaces - [x] Security report generation -- [ ] Extend to STRIDE, PASTA, LINDDUN, etc. +- [x] STRIDE threat modeling support +- [ ] Extend to PASTA, LINDDUN, etc. - [ ] Agent simulation-based threat exploration ## 🤝 Contributing diff --git a/src/repello_agent_wiz/analyzers/__init__.py b/src/repello_agent_wiz/analyzers/__init__.py index 8f4c672..dc17e9e 100644 --- a/src/repello_agent_wiz/analyzers/__init__.py +++ b/src/repello_agent_wiz/analyzers/__init__.py @@ -1 +1 @@ -from .threat_modelling import generate_maestro_analysis_report \ No newline at end of file +from .threat_modelling import generate_analysis_report, generate_maestro_analysis_report diff --git a/src/repello_agent_wiz/analyzers/stride.txt b/src/repello_agent_wiz/analyzers/stride.txt new file mode 100644 index 0000000..6df9a66 --- /dev/null +++ b/src/repello_agent_wiz/analyzers/stride.txt @@ -0,0 +1,98 @@ +STRIDE is a foundational threat modeling framework that categorizes adversarial goals into six classes—Spoofing, Tampering, Repudiation, Information Disclosure, Denial of Service, and Elevation of Privilege. +Applying STRIDE to agentic/AI-driven systems provides a structured way to examine how agents, tools, data, and communication channels might be abused or disrupted. + +1. Spoofing (Identity Misrepresentation) +Description: Attackers impersonate legitimate entities such as agents, users, APIs, or services to gain unauthorized access, issue commands, or exfiltrate data. In agentic systems, this often means forging inter-agent messages or pretending to be a trusted tool. +Example Scenarios: +- Malicious actor mimics a system agent to trigger destructive actions. +- Injected credentials chain into restricted tools or services. +- Prompt injection convinces an agent that hostile inputs are trusted instructions. +Potential Impacts: +- Unauthorized control of workflows or resources. +- Degradation of inter-agent trust and reliability. +- Leakage or corruption of sensitive information. +Mitigations: +- Strong authentication/authorization for all actors (human and non-human). +- Signed requests with cryptographic guarantees and regular key rotation. +- Secure identity attestation plus mutual TLS or scoped API keys. + +2. Tampering (Unauthorized Modification) +Description: Adversaries modify code, prompts, intermediate artifacts, or datasets to manipulate results. For AI agents this includes model poisoning, prompt alteration, or interference in data pipelines. +Example Scenarios: +- Malicious code inserted into tool plugins or runtime modules. +- Runtime manipulation of an agent’s working memory or system prompt. +- Poisoned embeddings, dataset labels, or configuration files. +Potential Impacts: +- Biased, erroneous, or unstable outputs. +- Downstream model/data corruption and loss of integrity. +- Difficult-to-trace behavioural drift. +Mitigations: +- Hash/signature validation, Merkle trees, and change detection. +- Immutable logging and version control for prompts, code, and data. +- Sandboxed execution, strict input validation, dependency attestation. + +3. Repudiation (Action Denial or Obfuscation) +Description: Attackers exploit weak auditability to deny or hide malicious actions. Distributed agent systems need transparent logs to ensure accountability. +Example Scenarios: +- Clearing or tampering with audit records after running unauthorized prompts. +- Omitted logging of critical tool invocations or parameter changes. +- Manipulated timestamps/metadata that mislead investigations. +Potential Impacts: +- Inability to attribute actions or detect insider threats. +- Delayed incident response and incomplete forensics. +- Erosion of trust in automated decisions. +Mitigations: +- Comprehensive, structured logging of agent reasoning and tool calls. +- Append-only/tamper-evident audit trails with cryptographic signing. +- Integration with SIEM tooling for correlation and anomaly detection. + +4. Information Disclosure (Data Exposure) +Description: Sensitive information is exposed through insecure storage, model outputs, or overly broad context sharing. +Example Scenarios: +- Agents revealing API keys or internal documents in responses. +- Model outputs leaking PII or proprietary training data. +- Tool integrations transmitting more context than necessary. +Potential Impacts: +- Regulatory compliance violations (GDPR, HIPAA, etc.). +- Intellectual property loss or competitive disadvantage. +- Amplified prompt-injection and model-inversion risks. +Mitigations: +- Data minimization with least-privilege context exposure. +- Encryption in transit/at rest plus strict access controls. +- Context redaction, confidential computing, output filtering. + +5. Denial of Service (Availability Disruption) +Description: Attackers exhaust resources or disrupt dependencies to halt agentic workflows. +Example Scenarios: +- Flooding an agent endpoint with traffic to exhaust capacity. +- Prompt loops that force agents into runaway processing. +- Cascading timeouts caused by dependency throttling or outages. +Potential Impacts: +- Workflow paralysis and downtime. +- Coordination breakdown across agents/tools. +- Increased operational costs and resource wastage. +Mitigations: +- Rate limiting, circuit breakers, and quota enforcement. +- Graceful degradation, redundancy, caching, failover plans. +- Continuous performance monitoring and anomaly detection. + +6. Elevation of Privilege (Unauthorized Capability Escalation) +Description: Adversaries gain higher privileges by chaining other STRIDE vectors or exploiting misconfigurations. +Example Scenarios: +- Compromised low-privilege agent invoking administrative tools. +- Prompt chaining that bypasses policy constraints. +- Misconfigured environment variables exposing critical secrets. +Potential Impacts: +- Full control of the system or persistent backdoors. +- Unauthorized data access and manipulation of workflows. +- Violation of least privilege and isolation assumptions. +Mitigations: +- Principle of Least Privilege and fine-grained RBAC. +- Segregation of duties between reasoning, execution, and control agents. +- Privilege audits, configuration drift detection, automated revocation. +Applying STRIDE to Agentic Workflows +- Asset & Flow Mapping: Enumerate agents, tools, data stores, channels, and external dependencies. +- Threat Enumeration: For each STRIDE category, identify plausible attack paths and adversarial behaviours. +- Impact & Likelihood Assessment: Score threats to prioritize remediation. +- Mitigation Planning: Develop layered controls spanning authn/z, isolation, observability, resilience. +- Validation & Monitoring: Continuously test defences and monitor for regressions or emerging threats. diff --git a/src/repello_agent_wiz/analyzers/sys_prompt.txt b/src/repello_agent_wiz/analyzers/sys_prompt_maestro.txt similarity index 100% rename from src/repello_agent_wiz/analyzers/sys_prompt.txt rename to src/repello_agent_wiz/analyzers/sys_prompt_maestro.txt diff --git a/src/repello_agent_wiz/analyzers/sys_prompt_stride.txt b/src/repello_agent_wiz/analyzers/sys_prompt_stride.txt new file mode 100644 index 0000000..f6868ef --- /dev/null +++ b/src/repello_agent_wiz/analyzers/sys_prompt_stride.txt @@ -0,0 +1,38 @@ +You are an expert in cybersecurity threat modeling. + +Your task is to perform a STRIDE-based threat analysis on the following agentic workflow graph. + +--- + +STRIDE Reference: + + +--- + +Agentic Workflow JSON Graph: + + +--- + +Produce a structured STRIDE threat modeling report in Markdown with these sections: + +STRIDE Analysis of Agentic Workflow +1. Mission Overview +Summarize what the system is trying to accomplish and the primary agent roles. Use short paragraphs. + +2. Asset & Data Inventory +Bullet the critical agents, tools, external systems, and sensitive data the workflow touches. + +3. STRIDE Threat Assessment +Create a table with columns: Category, Threat Description, Likelihood (Low/Medium/High), Impact (Low/Medium/High), Recommended Mitigations. Provide at least one row for every STRIDE category; add more if needed. + +4. Detailed Findings by Category +For each STRIDE category (Spoofing, Tampering, Repudiation, Information Disclosure, Denial of Service, Elevation of Privilege) write a short narrative describing how the threat manifests in this workflow, the potential consequences, and key mitigations. Reference specific agents/tools when possible. + +5. Prioritized Recommendations +List the top mitigations or design improvements the team should implement next, ordered by importance. + +Formatting requirements: +- Use valid Markdown headings and tables. +- Do not include any prose outside the report sections. +- Keep tone professional and action oriented. diff --git a/src/repello_agent_wiz/analyzers/threat_modelling.py b/src/repello_agent_wiz/analyzers/threat_modelling.py index 39d7d2e..729ee8a 100644 --- a/src/repello_agent_wiz/analyzers/threat_modelling.py +++ b/src/repello_agent_wiz/analyzers/threat_modelling.py @@ -1,5 +1,6 @@ import os import json +from typing import Dict from openai import OpenAI import importlib.resources as pkg_resources @@ -7,42 +8,64 @@ from repello_agent_wiz import analyzers -def generate_maestro_analysis_report(json_path: str): - # Load embedded files - with pkg_resources.files(analyzers).joinpath("maestro.txt").open("r") as f: - maestro = f.read() +_METHODOLOGY_CONFIGS: Dict[str, Dict[str, str]] = { + "maestro": { + "framework_file": "maestro.txt", + "prompt_template": "sys_prompt.txt", + "placeholder": "", + "label": "MAESTRO", + }, + "stride": { + "framework_file": "stride.txt", + "prompt_template": "sys_prompt_stride.txt", + "placeholder": "", + "label": "STRIDE", + }, +} - with pkg_resources.files(analyzers).joinpath("sys_prompt.txt").open("r") as f: - sys_prompt_template = f.read() - with open(json_path, "r") as f: - graph_data = json.load(f) - framework = graph_data.get("metadata", {}).get("framework", "unknown") +def generate_analysis_report(json_path: str, methodology: str = "maestro"): + method_key = methodology.lower() + if method_key not in _METHODOLOGY_CONFIGS: + valid = ", ".join(sorted(_METHODOLOGY_CONFIGS)) + raise ValueError(f"Unsupported methodology '{methodology}'. Valid options: {valid}") + + config = _METHODOLOGY_CONFIGS[method_key] + + # Load embedded reference material and prompt template + with pkg_resources.files(analyzers).joinpath(config["framework_file"]).open("r", encoding="utf-8") as f: + framework_reference = f.read() + + with pkg_resources.files(analyzers).joinpath(config["prompt_template"]).open("r", encoding="utf-8") as f: + sys_prompt_template = f.read() - with open(json_path, "r") as f: + with open(json_path, "r", encoding="utf-8") as f: graph_json = f.read() + graph_data = json.loads(graph_json) + framework = graph_data.get("metadata", {}).get("framework", "unknown") - sys_prompt = sys_prompt_template.replace("", maestro) + sys_prompt = sys_prompt_template.replace(config["placeholder"], framework_reference) sys_prompt = sys_prompt.replace("", graph_json) - # Initialize the OpenAI client properly client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) - # Use the client instance to create the completion response = client.chat.completions.create( model="gpt-4o", messages=[{"role": "user", "content": sys_prompt}], - temperature=0.3 + temperature=0.3, ) - # Extract content and remove markdown code block if present report = response.choices[0].message.content.strip() if report.startswith("```") and report.endswith("```"): report = "\n".join(report.splitlines()[1:-1]).strip() - - output_path = f"{framework}_report.md" - with open(output_path, "w") as f: + output_path = f"{framework}_{method_key}_report.md" + + with open(output_path, "w", encoding="utf-8") as f: f.write(report) - print(f"[✓] Saved MAESTRO analysis to: {output_path}") + print(f"[✓] Saved {config['label']} analysis to: {output_path}") + + +def generate_maestro_analysis_report(json_path: str): + generate_analysis_report(json_path, methodology="maestro") diff --git a/src/repello_agent_wiz/cli.py b/src/repello_agent_wiz/cli.py index fb1043f..7f01acd 100644 --- a/src/repello_agent_wiz/cli.py +++ b/src/repello_agent_wiz/cli.py @@ -1,6 +1,6 @@ import argparse from .frameworks import agent_chat, autogen, crewai, google_adk, langgraph, llama_index, n8n, openai_agents, pydantic, swarm -from .analyzers import generate_maestro_analysis_report +from .analyzers import generate_analysis_report from .visualizers.visualizer import generate_visualization @@ -20,6 +20,13 @@ def main(): # --- Analyze command --- analyze_parser = subparsers.add_parser("analyze", help="Run threat modeling analysis on extracted graph") analyze_parser.add_argument("--input", "-i", required=True, help="Path to JSON graph file") + analyze_parser.add_argument( + "--methodology", + "-m", + choices=["maestro", "stride"], + default="maestro", + help="Threat modeling framework to apply", + ) # --- Visualize command --- visualize_parser = subparsers.add_parser("visualize", help="Generate HTML visualization from graph JSON") @@ -57,7 +64,7 @@ def main(): print(f"Unknown framework: {args.framework}") case "analyze": - generate_maestro_analysis_report(args.input) + generate_analysis_report(args.input, methodology=args.methodology) case "visualize": generate_visualization(args.input, open_browser=args.open)