iflow-mcp-m507_ai-soc-agent 1.0.0__py3-none-any.whl

src/mcp/runbook_manager.py

@@ -0,0 +1,264 @@
+"""
+Runbook Manager for reading and parsing investigation runbooks.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+from ..core.errors import IntegrationError
+
+
+class RunbookManager:
+    """Manages runbook discovery and parsing."""
+    
+    def __init__(self, runbooks_dir: Optional[str] = None):
+        """
+        Initialize runbook manager.
+        
+        Args:
+            runbooks_dir: Path to runbooks directory.
+        """
+        if runbooks_dir is None:
+            # Default to run_books/ relative to project root
+            project_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+            runbooks_dir = os.path.join(project_root, "run_books")
+        
+        self.runbooks_dir = runbooks_dir
+    
+    def find_runbook(self, runbook_name: str, soc_tier: Optional[str] = None) -> Optional[str]:
+        """
+        Find runbook file by name.
+        
+        Args:
+            runbook_name: Name of runbook (e.g., "initial_alert_triage" or "soc1/triage/initial_alert_triage")
+            soc_tier: Optional SOC tier to limit search
+        
+        Returns:
+            Path to runbook file, or None if not found
+        """
+        # If runbook_name already includes path, use it directly
+        if "/" in runbook_name or runbook_name.startswith("soc"):
+            # Try as-is first
+            runbook_path = os.path.join(self.runbooks_dir, f"{runbook_name}.md")
+            if os.path.exists(runbook_path):
+                return runbook_path
+            
+            # Try without .md extension
+            runbook_path = os.path.join(self.runbooks_dir, runbook_name)
+            if os.path.exists(runbook_path):
+                return runbook_path
+        
+        # Search for runbook
+        for root, dirs, files in os.walk(self.runbooks_dir):
+            # Filter by soc_tier if provided
+            if soc_tier and f"/{soc_tier}/" not in root:
+                continue
+            
+            for file in files:
+                if file.endswith(".md"):
+                    # Check if filename matches (without extension)
+                    if file[:-3] == runbook_name or file[:-3].endswith(f"/{runbook_name}"):
+                        return os.path.join(root, file)
+                    
+                    # Check if filename contains runbook_name
+                    if runbook_name in file[:-3]:
+                        return os.path.join(root, file)
+        
+        return None
+    
+    def list_runbooks(
+        self,
+        soc_tier: Optional[str] = None,
+        category: Optional[str] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        List available runbooks.
+        
+        Args:
+            soc_tier: Filter by SOC tier (soc1, soc2, soc3)
+            category: Filter by category (triage, investigation, response, forensics, correlation)
+        
+        Returns:
+            List of runbook metadata dictionaries
+        """
+        runbooks = []
+        
+        if not os.path.exists(self.runbooks_dir):
+            return runbooks
+        
+        for root, dirs, files in os.walk(self.runbooks_dir):
+            # Filter by soc_tier
+            if soc_tier and f"/{soc_tier}/" not in root:
+                continue
+            
+            # Filter by category
+            if category:
+                if f"/{category}/" not in root:
+                    continue
+            
+            for file in files:
+                if file.endswith(".md") and file not in [
+                    "README.md",
+                    "index.md",
+                    "SOC_TIER_ORGANIZATION_PLAN.md",
+                    "IMPLEMENTATION_SUMMARY.md",
+                    "RUNBOOK_INTEGRATION_PROPOSAL.md",
+                    "AGENT_PROFILES_IMPLEMENTATION.md",
+                    "guidelines.md",  # SOC tier guidelines are not executable runbooks
+                ]:
+                    runbook_path = os.path.join(root, file)
+                    runbook_meta = self.parse_runbook_metadata(runbook_path)
+                    
+                    # Get relative path from runbooks_dir
+                    rel_path = os.path.relpath(runbook_path, self.runbooks_dir)
+                    runbook_name = rel_path[:-3]  # Remove .md extension
+                    
+                    runbooks.append({
+                        "name": runbook_name,
+                        "path": runbook_path,
+                        "soc_tier": runbook_meta.get("soc_tier"),
+                        "category": runbook_meta.get("category"),
+                        "objective": runbook_meta.get("objective", "")[:200],
+                        "description": runbook_meta.get("description", "")[:200]
+                    })
+        
+        return runbooks
+    
+    def read_runbook(self, runbook_path: str) -> str:
+        """
+        Read runbook content from file.
+        
+        Args:
+            runbook_path: Path to runbook file
+        
+        Returns:
+            Runbook content as string
+        """
+        if not os.path.exists(runbook_path):
+            raise IntegrationError(f"Runbook not found: {runbook_path}")
+        
+        with open(runbook_path, "r", encoding="utf-8") as f:
+            return f.read()
+    
+    def parse_runbook_metadata(self, runbook_path: str, content: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Parse metadata from runbook markdown.
+        
+        Args:
+            runbook_path: Path to runbook file
+            content: Optional runbook content (if already loaded)
+        
+        Returns:
+            Dictionary with parsed metadata
+        """
+        if content is None:
+            content = self.read_runbook(runbook_path)
+        
+        metadata = {}
+        
+        # Extract SOC tier from path
+        if "/soc1/" in runbook_path:
+            metadata["soc_tier"] = "soc1"
+        elif "/soc2/" in runbook_path:
+            metadata["soc_tier"] = "soc2"
+        elif "/soc3/" in runbook_path:
+            metadata["soc_tier"] = "soc3"
+        
+        # Extract category from path
+        if "/triage/" in runbook_path:
+            metadata["category"] = "triage"
+        elif "/investigation/" in runbook_path:
+            metadata["category"] = "investigation"
+        elif "/response/" in runbook_path:
+            metadata["category"] = "response"
+        elif "/forensics/" in runbook_path:
+            metadata["category"] = "forensics"
+        elif "/correlation/" in runbook_path:
+            metadata["category"] = "correlation"
+        elif "/enrichment/" in runbook_path:
+            metadata["category"] = "enrichment"
+        elif "/remediation/" in runbook_path:
+            metadata["category"] = "remediation"
+        elif "/cases/" in runbook_path:
+            # Case-specific runbooks are sub-runbooks
+            metadata["category"] = "cases"
+        
+        # Extract objective
+        obj_match = re.search(r"## Objective\s*\n\s*\n(.*?)(?=\n##|\Z)", content, re.DOTALL | re.IGNORECASE)
+        if obj_match:
+            metadata["objective"] = obj_match.group(1).strip()
+        
+        # Extract scope
+        scope_match = re.search(r"## Scope\s*\n\s*\n(.*?)(?=\n##|\Z)", content, re.DOTALL | re.IGNORECASE)
+        if scope_match:
+            metadata["scope"] = scope_match.group(1).strip()
+        
+        # Extract tools
+        tools_match = re.search(r"## Tools\s*\n\s*\n(.*?)(?=\n##|\Z)", content, re.DOTALL | re.IGNORECASE)
+        if tools_match:
+            tools_text = tools_match.group(1)
+            # Extract tool names (look for backtick-wrapped tool names)
+            tool_names = re.findall(r"`([a-z_]+)`", tools_text, re.IGNORECASE)
+            metadata["tools"] = list(set(tool_names))  # Remove duplicates
+        
+        # Extract inputs
+        inputs_match = re.search(r"## Inputs\s*\n\s*\n(.*?)(?=\n##|\Z)", content, re.DOTALL | re.IGNORECASE)
+        if inputs_match:
+            inputs_text = inputs_match.group(1)
+            # Extract input variables
+            input_vars = re.findall(r"\$\{([A-Z_]+)\}", inputs_text)
+            metadata["inputs"] = list(set(input_vars))
+        
+        # Extract workflow steps count
+        steps_match = re.search(r"## Workflow Steps", content, re.IGNORECASE)
+        if steps_match:
+            # Count numbered steps
+            step_count = len(re.findall(r"^\d+\.\s+\*\*", content, re.MULTILINE))
+            metadata["step_count"] = step_count
+        
+        return metadata
+    
+    def extract_workflow_steps(self, content: str) -> List[Dict[str, Any]]:
+        """
+        Extract workflow steps from runbook content.
+        
+        Args:
+            content: Runbook content
+        
+        Returns:
+            List of workflow step dictionaries
+        """
+        steps = []
+        
+        # Find workflow steps section
+        steps_section_match = re.search(
+            r"## Workflow Steps.*?\n(.*?)(?=\n```|\n##|\Z)",
+            content,
+            re.DOTALL | re.IGNORECASE
+        )
+        
+        if not steps_section_match:
+            return steps
+        
+        steps_text = steps_section_match.group(1)
+        
+        # Extract numbered steps
+        step_pattern = r"(\d+)\.\s+\*\*(.*?)\*\*:(.*?)(?=\d+\.\s+\*\*|\Z)"
+        step_matches = re.finditer(step_pattern, steps_text, re.DOTALL)
+        
+        for match in step_matches:
+            step_num = int(match.group(1))
+            step_title = match.group(2).strip()
+            step_content = match.group(3).strip()
+            
+            steps.append({
+                "step_number": step_num,
+                "title": step_title,
+                "content": step_content
+            })
+        
+        return steps
+

src/orchestrator/__init__.py

@@ -0,0 +1,11 @@
+"""
+Orchestrator and LLM tools for SamiGPT.
+
+This package contains:
+- ``incident_workflow.py``: High-level workflows that coordinate between
+  case management, SIEM, and EDR clients.
+- ``tools_case.py``: LLM-callable tools for case management operations.
+- ``tools_siem.py``: LLM-callable tools for SIEM operations.
+- ``tools_edr.py``: LLM-callable tools for EDR operations.
+"""
+

src/orchestrator/incident_workflow.py

@@ -0,0 +1,244 @@
+"""
+High-level incident response workflows for SamiGPT.
+
+This module provides orchestration functions that coordinate between
+case management, SIEM, and EDR clients to perform common incident
+response tasks.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from ..api.case_management import (
+    Case,
+    CaseManagementClient,
+    CaseObservable,
+    CaseStatus,
+)
+from ..api.edr import EDRClient, Endpoint
+from ..api.siem import SIEMClient, SiemAlert
+from ..core.errors import IntegrationError
+
+
+def create_incident_from_alert(
+    alert: SiemAlert,
+    case_client: CaseManagementClient,
+    title_prefix: Optional[str] = None,
+) -> Case:
+    """
+    Create a new case in the case management system from a SIEM alert.
+
+    Args:
+        alert: The SIEM alert to convert into a case.
+        case_client: The case management client to use.
+        title_prefix: Optional prefix for the case title.
+
+    Returns:
+        The created case.
+
+    Raises:
+        IntegrationError: If case creation fails.
+    """
+    title = f"{title_prefix + ': ' if title_prefix else ''}{alert.title or 'Security Alert'}"
+    description = f"Alert from SIEM: {alert.description or 'No description'}\n\n"
+    description += f"Severity: {alert.severity}\n"
+    description += f"Source: {alert.source}\n"
+    if alert.timestamp:
+        description += f"Timestamp: {alert.timestamp.isoformat()}\n"
+
+    case = Case(
+        title=title,
+        description=description,
+        status=CaseStatus.OPEN,
+        priority=alert.severity.value if hasattr(alert.severity, "value") else "medium",
+        tags=alert.tags or [],
+    )
+
+    try:
+        created_case = case_client.create_case(case)
+        return created_case
+    except Exception as e:
+        raise IntegrationError(f"Failed to create case from alert: {e}") from e
+
+
+def enrich_case_from_siem(
+    case_id: str,
+    case_client: CaseManagementClient,
+    siem_client: SIEMClient,
+    search_terms: Optional[List[str]] = None,
+) -> List[CaseObservable]:
+    """
+    Enrich a case by searching SIEM for related events and adding observables.
+
+    Args:
+        case_id: The ID of the case to enrich.
+        case_client: The case management client.
+        siem_client: The SIEM client to search.
+        search_terms: Optional list of terms to search for in SIEM.
+
+    Returns:
+        List of observables added to the case.
+
+    Raises:
+        IntegrationError: If enrichment fails.
+    """
+    try:
+        case = case_client.get_case(case_id)
+    except Exception as e:
+        raise IntegrationError(f"Failed to retrieve case {case_id}: {e}") from e
+
+    observables_added = []
+
+    # Extract potential observables from case title/description
+    search_terms = search_terms or []
+    if case.title:
+        search_terms.append(case.title)
+    if case.description:
+        # Simple extraction - in production, use more sophisticated parsing
+        search_terms.append(case.description[:100])
+
+    # Search SIEM for related events
+    try:
+        if search_terms:
+            query = " OR ".join(search_terms[:5])  # Limit to avoid huge queries
+            events = siem_client.search_security_events(
+                query=query,
+                limit=50,
+            )
+
+            # Extract unique observables from events
+            seen_observables = set()
+            for event in events.events:
+                # Add IP addresses as observables
+                if event.ip:
+                    key = f"ip:{event.ip}"
+                    if key not in seen_observables:
+                        observable = CaseObservable(
+                            type="ip",
+                            value=event.ip,
+                            description=f"Found in SIEM event: {event.id}",
+                        )
+                        case_client.add_case_observable(case_id, observable)
+                        observables_added.append(observable)
+                        seen_observables.add(key)
+
+                # Add file hashes if present
+                if event.file_hash:
+                    key = f"hash:{event.file_hash}"
+                    if key not in seen_observables:
+                        observable = CaseObservable(
+                            type="hash",
+                            value=event.file_hash,
+                            description=f"Found in SIEM event: {event.id}",
+                        )
+                        case_client.add_case_observable(case_id, observable)
+                        observables_added.append(observable)
+                        seen_observables.add(key)
+
+    except Exception as e:
+        raise IntegrationError(f"Failed to search SIEM or add observables: {e}") from e
+
+    return observables_added
+
+
+def enrich_case_from_edr(
+    case_id: str,
+    case_client: CaseManagementClient,
+    edr_client: EDRClient,
+    endpoint_id: Optional[str] = None,
+) -> List[CaseObservable]:
+    """
+    Enrich a case by querying EDR for endpoint information and adding observables.
+
+    Args:
+        case_id: The ID of the case to enrich.
+        case_client: The case management client.
+        edr_client: The EDR client to query.
+        endpoint_id: Optional endpoint ID to focus on. If not provided, uses
+            observables from the case to find relevant endpoints.
+
+    Returns:
+        List of observables added to the case.
+
+    Raises:
+        IntegrationError: If enrichment fails.
+    """
+    try:
+        case = case_client.get_case(case_id)
+    except Exception as e:
+        raise IntegrationError(f"Failed to retrieve case {case_id}: {e}") from e
+
+    observables_added = []
+
+    try:
+        # If endpoint_id provided, get that endpoint's details
+        if endpoint_id:
+            endpoint = edr_client.get_endpoint_summary(endpoint_id)
+            if endpoint:
+                # Add endpoint hostname as observable
+                if endpoint.hostname:
+                    observable = CaseObservable(
+                        type="hostname",
+                        value=endpoint.hostname,
+                        description=f"Endpoint from EDR: {endpoint_id}",
+                    )
+                    case_client.add_case_observable(case_id, observable)
+                    observables_added.append(observable)
+
+        # Get recent detections and add file hashes as observables
+        detections = edr_client.list_detections(limit=20)
+        seen_hashes = set()
+        for detection in detections:
+            if detection.file_hash:
+                key = f"hash:{detection.file_hash}"
+                if key not in seen_hashes:
+                    observable = CaseObservable(
+                        type="hash",
+                        value=detection.file_hash,
+                        description=f"EDR detection: {detection.id}",
+                    )
+                    case_client.add_case_observable(case_id, observable)
+                    observables_added.append(observable)
+                    seen_hashes.add(key)
+
+    except Exception as e:
+        raise IntegrationError(f"Failed to query EDR or add observables: {e}") from e
+
+    return observables_added
+
+
+def close_incident(
+    case_id: str,
+    case_client: CaseManagementClient,
+    resolution_notes: Optional[str] = None,
+) -> Case:
+    """
+    Close an incident case with optional resolution notes.
+
+    Args:
+        case_id: The ID of the case to close.
+        case_client: The case management client.
+        resolution_notes: Optional notes about the resolution.
+
+    Returns:
+        The updated case.
+
+    Raises:
+        IntegrationError: If closing the case fails.
+    """
+    try:
+        # Add resolution notes as a comment if provided
+        if resolution_notes:
+            case_client.add_case_comment(
+                case_id=case_id,
+                content=f"Resolution: {resolution_notes}",
+                author=None,  # System-generated
+            )
+
+        # Update status to closed
+        updated_case = case_client.update_case_status(case_id, CaseStatus.CLOSED)
+        return updated_case
+    except Exception as e:
+        raise IntegrationError(f"Failed to close case {case_id}: {e}") from e
+