diff --git a/cyberai/agents/exploit/agent.py b/cyberai/agents/exploit/agent.py index 8a79c22..7e4ae1d 100644 --- a/cyberai/agents/exploit/agent.py +++ b/cyberai/agents/exploit/agent.py @@ -14,6 +14,11 @@ from .attack_path import AttackPath, build_attack_paths from .chain_builder import build_exploit_chain +from cyberai.core.types import ( + AttackPath as AttackPathModel, + ExploitChain as ExploitChainModel, + ExploitResult, +) from .cvss_analyzer import analyze_attack_vector console = Console() @@ -69,6 +74,21 @@ def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str "exploit_chain": chain, "ai_analysis": analysis, } + # Build a validated ExploitResult and store it in the KB. + exploit_result = ExploitResult( + target=target, + attack_paths=[ + AttackPathModel(**p.to_dict()) for p in attack_paths + ], + chain=( + ExploitChainModel(**chain) + if isinstance(chain, dict) else None + ), + ai_analysis=analysis, + ) + self.kb.set( + "exploit.result", exploit_result.model_dump(), agent=self.AGENT_NAME + ) self.kb.set("exploit", result, agent=self.AGENT_NAME) self._log("Exploit analysis complete", {"paths_found": len(attack_paths)}) return result diff --git a/cyberai/agents/intel/agent.py b/cyberai/agents/intel/agent.py index de19bf3..6b942d4 100644 --- a/cyberai/agents/intel/agent.py +++ b/cyberai/agents/intel/agent.py @@ -9,6 +9,7 @@ from .nvd_client import get_cve, search_cves from .service_mapper import ports_to_queries, score_to_severity +from cyberai.core.types import CVEEntry, IntelResult class IntelAgent(BaseAgent): @@ -80,6 +81,30 @@ def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str (cve.get("cvss", {}) or {}).get("vector", "")], ) + # Build a validated IntelResult and store it in the KB. + def _cvss_score(c: dict) -> float: + raw = c.get("cvss") + if isinstance(raw, dict): + return float(raw.get("score") or 0.0) + return float(raw or 0.0) + + intel_result = IntelResult( + target=target, + cves=[ + CVEEntry( + id=c.get("id") or c.get("cve_id", ""), + cvss=_cvss_score(c), + severity=score_to_severity(_cvss_score(c)), + description=c.get("description", ""), + published=c.get("published") or None, + exploited_in_wild=c.get("exploited_in_wild", False), + epss=float(c.get("epss") or 0.0), + ) + for c in all_cves + ], + ) + self.kb.set("intel.result", intel_result.model_dump(), agent=self.AGENT_NAME) + result = { "status": "done", "queries": queries, diff --git a/cyberai/agents/recon/agent.py b/cyberai/agents/recon/agent.py index 1a88ace..a004b5d 100644 --- a/cyberai/agents/recon/agent.py +++ b/cyberai/agents/recon/agent.py @@ -5,6 +5,7 @@ from cyberai.core.base_agent import BaseAgent, Tool from cyberai.core.scan_session import Severity +from cyberai.core.types import OpenPort, ReconResult from .dns_tool import detect_subdomains, run_dns, run_whois from .nmap_tool import run_nmap @@ -88,4 +89,17 @@ def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str evidence=[str(p) for p in ports], ) + # Build a validated pydantic ReconResult and store it in the KB. + recon_result = ReconResult( + target=target, + ports=[OpenPort(**p) for p in ports if isinstance(p, dict)], + whois=whois_result if isinstance(whois_result, dict) else {}, + dns=dns_result if isinstance(dns_result, dict) else {}, + subdomains=( + sub_result.get("subdomains", []) + if isinstance(sub_result, dict) else [] + ), + ) + self.kb.set("recon.result", recon_result.model_dump(), agent=self.AGENT_NAME) + return {"status": "done", "kb_keys": list(results.keys()), "ports": ports} diff --git a/cyberai/core/types.py b/cyberai/core/types.py index 0b422f6..656f7f4 100644 --- a/cyberai/core/types.py +++ b/cyberai/core/types.py @@ -5,6 +5,8 @@ from typing import Any, Union from pathlib import Path +from pydantic import BaseModel, Field + # Target types Target = str # IP, CIDR, or domain PortNumber = int # 1-65535 @@ -16,18 +18,77 @@ # Recon PortList = list[PortNumber] -ServiceMap = dict[str, ServiceName] # "80" → "http" -ReconResult = dict[str, Any] +ServiceMap = dict[str, ServiceName] # "80" -> "http" + + +class OpenPort(BaseModel): + """A single open port discovered during recon.""" + port: int + protocol: str = "tcp" + service: str = "unknown" + version: str | None = None + + +class ReconResult(BaseModel): + """Structured output of the ReconAgent.""" + target: str + ports: list[OpenPort] = Field(default_factory=list) + whois: dict[str, Any] = Field(default_factory=dict) + dns: dict[str, Any] = Field(default_factory=dict) + subdomains: list[str] = Field(default_factory=list) # Intel CVEId = str # "CVE-2024-1234" CVEList = list[CVEId] RiskLevel = str # "CRITICAL" | "HIGH" | "MEDIUM" | "LOW" -IntelResult = dict[str, Any] + + +class CVEEntry(BaseModel): + """A single CVE with scoring and threat-intel context.""" + id: str + cvss: float = 0.0 + severity: str = "UNKNOWN" + description: str = "" + published: str | None = None + exploited_in_wild: bool = False + epss: float = 0.0 + + +class IntelResult(BaseModel): + """Structured output of the IntelAgent.""" + target: str + cves: list[CVEEntry] = Field(default_factory=list) + services: dict[str, Any] = Field(default_factory=dict) # Exploit -AttackPath = dict[str, Any] -ExploitResult = dict[str, Any] + + +class AttackPath(BaseModel): + """A single attack path derived from one CVE.""" + cve_id: str + attack_vector: str = "Unknown" + attack_complexity: str = "Unknown" + technique: str = "" + success_probability: float = 0.0 + requires_auth: bool = False + requires_interaction: bool = False + notes: str = "" + + +class ExploitChain(BaseModel): + """An ordered, MITRE-mapped sequence of exploitation steps.""" + target: str + chain_length: int = 0 + steps: list[dict[str, Any]] = Field(default_factory=list) + summary: str = "" + + +class ExploitResult(BaseModel): + """Structured output of the ExploitAgent.""" + target: str + attack_paths: list[AttackPath] = Field(default_factory=list) + chain: ExploitChain | None = None + ai_analysis: str = "" # Report ReportPath = Path