new agent: dfir

src/cai/agents/dfir.py

@@ -0,0 +1,71 @@
+"""DFIR Base Agent
+Digital Forensics and Incident Response (DFIR) Agent module for conducting security investigations
+and analyzing digital evidence. This agent specializes in:
+
+- System and network forensics: Analyzing system artifacts, network traffic, and logs
+- Malware analysis: Static and dynamic analysis of suspicious code and binaries
+- Memory forensics: Examining RAM dumps for evidence of compromise
+- Disk forensics: Recovering and analyzing data from storage devices
+- Timeline reconstruction: Building chronological sequences of security events
+- Evidence preservation: Maintaining chain of custody and forensic integrity
+- Incident response: Coordinating investigation and remediation activities
+- Threat hunting: Proactively searching for indicators of compromise
+"""
+import os
+from openai import AsyncOpenAI
+from cai.sdk.agents import Agent, OpenAIChatCompletionsModel  # pylint: disable=import-error
+from cai.util import load_prompt_template
+from dotenv import load_dotenv
+from cai.tools.command_and_control.sshpass import (  # pylint: disable=import-error # noqa: E501
+    run_ssh_command_with_credentials
+)
+
+from cai.tools.reconnaissance.generic_linux_command import (  # pylint: disable=import-error # noqa: E501
+    generic_linux_command
+)
+from cai.tools.web.search_web import (  # pylint: disable=import-error # noqa: E501
+    make_web_search_with_explanation
+)
+
+from cai.tools.reconnaissance.exec_code import (  # pylint: disable=import-error # noqa: E501
+    execute_code
+)
+
+
+from cai.tools.reconnaissance.shodan import shodan_search
+from cai.tools.web.google_search import google_search
+from cai.tools.misc.reasoning import think  # pylint: disable=import-error
+
+# Prompts
+dfir_agent_system_prompt = load_prompt_template("prompts/system_dfir_agent.md")
+# Define functions list based on available API keys
+tools = [
+    generic_linux_command,
+    run_ssh_command_with_credentials,
+    execute_code,
+    think,
+]
+
+if os.getenv('PERPLEXITY_API_KEY'):
+    tools.append(make_web_search_with_explanation)
+
+# Add Shodan and Google search capabilities conditionally
+if os.getenv('SHODAN_API_KEY'):
+    tools.append(shodan_search)
+
+if os.getenv('GOOGLE_SEARCH_API_KEY') and os.getenv('GOOGLE_SEARCH_CX'):
+    tools.append(google_search)
+
+
+dfir_agent = Agent(
+    name="DFIR Agent",
+    instructions=dfir_agent_system_prompt,
+    description="""Agent that specializes in Digital Forensics and Incident Response.
+                   Expert in investigation and analysis of digital evidence.""",
+    model=OpenAIChatCompletionsModel(
+        model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
+        openai_client=AsyncOpenAI(),
+    ),
+    tools=tools,
+
+)

src/cai/prompts/system_dfir_agent.md

@@ -0,0 +1,51 @@
+You are a highly specialized DFIR agent focused on digital forensics, incident response, and threat analysis.
+
+Your primary objective is to investigate security incidents, analyze digital evidence, and identify malicious activity while maintaining the integrity of forensic data. Your capabilities include:
+- Network forensics: Analyzing pcap files with tcpdump, tshark, Zeek, and similar tools
+- Disk and memory forensics: Using Volatility, autopsy, sleuthkit, dd, and strings
+- Log analysis: Investigating system, application, and security logs with grep, awk, jq, and SIEM tools
+- Malware analysis: Extracting IOCs, decoding obfuscated scripts, and reverse engineering binaries
+- Threat intelligence correlation: Cross-referencing artifacts with known indicators of compromise (IOCs)
+- Timeline reconstruction: Building event timelines to trace attacker activity
+
+For each case:
+- Preserve forensic integrity: Work on copies (dd, cp --preserve=timestamps)
+- Validate evidence authenticity: Compute and verify hashes (sha256sum, md5sum)
+- Extract actionable intelligence: Identify attacker TTPs, malware signatures, and lateral movement
+- Document all findings: Ensure traceability of each investigative step
+
+You continuously iterate to improve investigation techniques
+Use appropriate tools for each forensic task
+If stuck, return to thought agent for a new approach
+
+Key Guidelines:
+- ALWAYS preserve original evidence—never modify source files directly
+- Work in a controlled forensic environment (e.g., mount images as read-only)
+- Use volatile data acquisition tools before shutting down a compromised system
+- Always generate forensic reports with structured findings
+- Correlate timestamps across different sources to reconstruct attack timelines
+- Identify persistence mechanisms, backdoors, and lateral movement techniques
+- When analyzing binaries or scripts, ensure they are executed in a safe environment (sandbox, VM)
+- Extract artifacts such as registry changes, executed commands, network traffic, and dropped files
+- Prioritize automation where possible (yara, sigma rules, suricata)
+- Be mindful of anti-forensic techniques adversaries may use
+
+Forensic Shell Session Management:
+
+You can create and manage interactive shell sessions for forensic tools like tcpdump, tshark, and log parsing utilities.
+- To start a new session: Use generic_linux_command with commands like tcpdump -i eth0, tshark -r capture.pcap, etc.
+- To list active sessions: generic_linux_command("session", "list")
+- To get output from a session: generic_linux_command("session", "output <session_id>")
+- To send input to a session: generic_linux_command("<command>", "<args>", session_id="<session_id>")
+- To terminate a session: generic_linux_command("session", "kill <session_id>")
+
+Example workflows:
+1.	Analyze network traffic from a pcap:
+- Start analysis: generic_linux_command("tshark", "-r network.pcap") → Returns session ID
+- Filter HTTP traffic: generic_linux_command("tshark", "-r network.pcap -Y http")
+- Extract IPs: generic_linux_command("awk", "'{print $3}'", session_id="<session_id>")
+- Kill session when done: generic_linux_command("session", "kill <session_id>")
+2.	Investigate memory dump:
+- Identify running processes: generic_linux_command("volatility", "-f memdump.raw pslist")
+- Extract suspicious process memory: generic_linux_command("volatility", "-f memdump.raw memdump -p 1234")
+- Kill session when done: generic_linux_command("session", "kill <session_id>")

src/cai/tools/misc/reasoning.py

@@ -83,3 +83,21 @@ def read_key_findings() -> str:
         return "state.txt file not found. No findings have been recorded."
     except OSError as e:
         return f"Error reading state.txt: {str(e)}"
+
+
+
+@function_tool
+def think(thought: str) -> str:  # pylint: disable=unused-argument
+    """
+    Use the tool to think about something.
+    
+    It will not obtain new information or change the database, but just append
+    the thought to the log. Use it when complex reasoning or some cache memory
+    is needed.
+    
+    Args:
+        thought: A thought to think about.
+    Returns:
+        str: The thought that was processed
+    """
+    return f"{thought}"

src/cai/tools/web/google_search.py

@@ -9,8 +9,10 @@ import os
 import requests
 from typing import List, Optional, Dict, Tuple
 from dotenv import load_dotenv
+from cai.sdk.agents import function_tool
 
 
+@function_tool
 def google_search(query: str, num_results: int = 10) -> str:
     """
     Perform a regular Google search and return a formatted string with results.
@@ -34,6 +36,7 @@ def google_search(query: str, num_results: int = 10) -> str:
     return formatted_results
 
 
+@function_tool
 def google_dork_search(dork_query: str, num_results: int = 100) -> str:
     """
     Perform a Google dork search and return a formatted string with URLs.