diff --git a/src/cai/agents/dfir.py b/src/cai/agents/dfir.py new file mode 100644 index 00000000..c2a4902e --- /dev/null +++ b/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, + +) \ No newline at end of file diff --git a/src/cai/prompts/system_dfir_agent.md b/src/cai/prompts/system_dfir_agent.md new file mode 100644 index 00000000..d4a3b2af --- /dev/null +++ b/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 ") +- To send input to a session: generic_linux_command("", "", session_id="") +- To terminate a session: generic_linux_command("session", "kill ") + +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="") +- Kill session when done: generic_linux_command("session", "kill ") +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 ") \ No newline at end of file diff --git a/src/cai/tools/misc/reasoning.py b/src/cai/tools/misc/reasoning.py index b793eca1..25e9684a 100644 --- a/src/cai/tools/misc/reasoning.py +++ b/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}" diff --git a/src/cai/tools/web/google_search.py b/src/cai/tools/web/google_search.py index c5c402df..26fd6be2 100644 --- a/src/cai/tools/web/google_search.py +++ b/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.