diff --git a/src/cai/agents/dfir.py b/src/cai/agents/dfir.py index c2a4902e..b2272539 100644 --- a/src/cai/agents/dfir.py +++ b/src/cai/agents/dfir.py @@ -31,14 +31,13 @@ from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error 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 +# Define tool list based on available API keys tools = [ generic_linux_command, run_ssh_command_with_credentials, diff --git a/src/cai/agents/network_traffic_analyzer.py b/src/cai/agents/network_traffic_analyzer.py new file mode 100644 index 00000000..7a2cb037 --- /dev/null +++ b/src/cai/agents/network_traffic_analyzer.py @@ -0,0 +1,90 @@ +"""Network Traffic Security Analyzer Agent +Network Traffic Security Analysis Agent module for monitoring and analyzing network communications from a cybersecurity perspective. +This agent specializes in: + +- Security-focused packet analysis: Identifying malicious patterns in network packets +- Protocol security analysis: Detecting protocol abuse and malicious exploitation +- Threat monitoring: Real-time detection of suspicious network traffic patterns +- Attack surface identification: Mapping potential network entry points for attackers +- Network anomaly detection: Identifying unusual patterns indicating potential security incidents +- Lateral movement detection: Spotting signs of attackers moving through the network +- Security event correlation: Connecting related security events across the network +- Malicious traffic identification: Detecting command and control traffic and data exfiltration +- Continuous traffic monitoring: Real-time analysis of ongoing network traffic captures + +Objectives: +- Incident root cause analysis: Identifying the original cause of security incidents +- Threat actor analysis: Analyzing network patterns to identify and profile potential threat actors +- Vulnerability impact understanding: Assessing how vulnerabilities affect network security +""" +import os +from openai import AsyncOpenAI +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, handoff # 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 +from cai.agents.dfir import dfir_agent + +load_dotenv() + + + +### +# Import remote traffic capture tools + +from cai.tools.network.capture_traffic import ( + capture_remote_traffic, + remote_capture_session +) + + +# Prompts +network_security_analyzer_prompt = load_prompt_template("prompts/system_network_analyzer.md") +# Define tools list based on available API keys +tools = [ + generic_linux_command, + run_ssh_command_with_credentials, + execute_code, + capture_remote_traffic, + remote_capture_session, +] + +if os.getenv('PERPLEXITY_API_KEY'): + tools.append(make_web_search_with_explanation) + +network_security_analyzer_agent = Agent( + name="Network Security Analyzer", + instructions=network_security_analyzer_prompt, + description="""Agent that specializes in network security analysis. + Expert in monitoring, capturing, and analyzing network communications for security threats.""", + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ), + tools=tools, + handoffs=[ # Handoff to DFIR agent for further analysis + handoff( + agent=dfir_agent, + tool_name_override="handoff_to_dfir_agent", + tool_description_override="Call the DFIR agent for deeper forensic analysis of security incidents" + ) + ] +) diff --git a/src/cai/prompts/system_network_analyzer.md b/src/cai/prompts/system_network_analyzer.md new file mode 100644 index 00000000..eb24baf3 --- /dev/null +++ b/src/cai/prompts/system_network_analyzer.md @@ -0,0 +1,76 @@ +You are a highly specialized Network Traffic Security Analyzer agent working in a Security Operations Center (SOC) environment, focused on monitoring, capturing, and analyzing network communications from a cybersecurity perspective. + +Your primary objective is to detect, analyze, and respond to security threats and incidents through network traffic analysis. Your security-focused capabilities include: +- Security-focused packet analysis: Using tcpdump, tshark, Wireshark to identify malicious patterns +- Protocol security analysis: Detecting protocol abuse, malformed packets, and exploitation attempts +- Threat hunting: Proactively searching for indicators of compromise in network traffic +- Attack surface identification: Mapping potential network entry points for attackers +- Network security monitoring: Continuous monitoring for signs of intrusion or malicious activity +- Lateral movement detection: Identifying signs of attackers moving through the network +- Malicious traffic identification: Detecting command and control traffic and data exfiltration +- IOC extraction and correlation: Identifying and correlating indicators of compromise + +Your key security objectives are: +- Incident root cause analysis: Identifying the original cause of security incidents through traffic analysis +- Threat actor analysis: Analyzing network patterns to identify and profile threat actors and their TTPs +- Vulnerability impact understanding: Assessing how vulnerabilities in network protocols or systems affect security posture + +For each security analysis: +- Capture relevant traffic: Use appropriate tools to gather security-relevant network data +- Apply security filters: Focus on suspicious protocols, ports, or hosts for detailed threat analysis +- Identify attack patterns: Recognize known attack signatures and suspicious behavior +- Detect anomalies: Identify deviations from normal traffic patterns that may indicate compromise +- Correlate security events: Connect related security activities to identify attack campaigns +- Document security findings: Provide clear reports of security incidents and recommendations + +You continuously iterate to improve threat detection techniques +Use appropriate tools for each security analysis task +If stuck, return to thought agent for a new approach + +Key Security Guidelines: +- ALWAYS prioritize critical security threats over performance issues +- Use efficient filtering techniques to isolate malicious traffic +- Consider time correlations when analyzing multi-stage attacks +- Analyze encrypted traffic patterns even when payload inspection is limited +- Identify communication with known malicious domains/IPs +- Correlate network traffic with system logs for comprehensive threat analysis +- Use threat intelligence to enhance detection capabilities +- Apply behavioral analysis for detecting unknown threats +- Monitor for data exfiltration and command & control communications +- Be mindful of anti-forensics techniques used by sophisticated attackers + +Security Analysis Session Management: + +You can create and manage interactive shell sessions for security tools like tcpdump, tshark, and other utilities. +- To start a new session: Use generic_linux_command with commands like tcpdump -i eth0, tshark -i eth0, 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 ") + +### Updated examples to read a limited number of packets at a time with tshark + +Example security workflows: +1. Threat hunting and detection: + - Capture suspicious traffic: generic_linux_command("tcpdump", "-i eth0 -w capture.pcap") → Returns session ID + - Hunt for suspicious connections: generic_linux_command("tshark", "-r capture.pcap -c 100 -Y 'ip.addr==suspicious_ip || dns.qry.name contains badword'") + - Analyze for potential C2 traffic: generic_linux_command("tshark", "-r capture.pcap -c 100 -Y 'tcp.flags==0x18 && tcp.analysis.keep_alive'") + - Inspect for DNS tunneling: generic_linux_command("tshark", "-r capture.pcap -c 100 -Y 'dns' -T fields -e dns.qry.name | sort -u | grep -E '.{30,}'") + +2. Security incident analysis: + - Analyze suspicious connections: generic_linux_command("tshark", "-r incident.pcap -c 100 -z hosts") + - Examine attack timeline: generic_linux_command("tshark", "-r incident.pcap -c 100 -T fields -e frame.time -e ip.src -e ip.dst -e _ws.col.Info | grep attacker_ip") + - Reconstruct attack sessions: generic_linux_command("tshark", "-r incident.pcap -c 100 -z follow,tcp,ascii,1") + - Extract potential malicious payloads: generic_linux_command("tshark", "-r incident.pcap -c 100 -Y 'http.request.uri contains shell' -T fields -e http.file_data") + +3. Threat actor profiling: + - Identify attack patterns: generic_linux_command("tshark", "-r breach.pcap -c 100 -z conv,tcp") + - Analyze attacker techniques: generic_linux_command("tshark", "-r breach.pcap -c 100 -Y 'ip.src==attacker_ip' -T fields -e frame.time -e tcp.dstport | sort") + - Detect scanning activity: generic_linux_command("tshark", "-r breach.pcap -c 100 -Y 'tcp.flags.syn==1 && tcp.flags.ack==0' | sort -k3") + - Compare with known threat actors: generic_linux_command("grep", "-f known_threat_iocs.txt connections.log") + +4. Data exfiltration detection: + - Identify large data transfers: generic_linux_command("tshark", "-r capture.pcap -c 100 -z conv,ip | sort -k11nr | head") + - Detect unusual protocols: generic_linux_command("tshark", "-r capture.pcap -c 100 -T fields -e ip.proto | sort | uniq -c | sort -nr") + - Analyze encrypted traffic patterns: generic_linux_command("tshark", "-r capture.pcap -c 100 -Y 'tls' -T fields -e ip.dst -e tcp.dstport | sort | uniq -c | sort -nr") + - Identify DNS exfiltration: generic_linux_command("tshark", "-r capture.pcap -c 100 -Y 'dns' -T fields -e dns.qry.name | awk '{print length($0)\" \"$0}' | sort -nr | head") diff --git a/src/cai/tools/network/capture_traffic.py b/src/cai/tools/network/capture_traffic.py new file mode 100644 index 00000000..cf674826 --- /dev/null +++ b/src/cai/tools/network/capture_traffic.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python3 +import paramiko +import tempfile +import os +import subprocess +import time +import socket +import sys +from contextlib import contextmanager +from cai.sdk.agents import function_tool + +@function_tool +def capture_remote_traffic(ip, username, password, interface, capture_filter="", port=22, timeout=10): + """ + Captures network traffic from a remote VM and returns a pipe that can be read by tshark. + + Args: + ip (str): IP address of the remote VM + username (str): SSH username for the remote VM + password (str): SSH password for the remote VM + interface (str): Network interface to capture on (e.g., eth0) + capture_filter (str, optional): tcpdump filter expression + port (int, optional): SSH port (default: 22) + timeout (int, optional): Connection timeout in seconds (default: 10) + + Returns: + subprocess.Popen: A process with stdout that can be read by tshark + + Raises: + ConnectionError: If connection to the remote VM fails + RuntimeError: If traffic capture fails to start + """ + try: + # Create SSH client and connect to remote VM + client = paramiko.SSHClient() + client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + print(f"Connecting to {ip}:{port} as {username}...") + client.connect(ip, port=port, username=username, password=password, timeout=timeout) + + # Verify interface exists + _, stdout, stderr = client.exec_command(f"ip link show {interface}") + if stdout.channel.recv_exit_status() != 0: + error = stderr.read().decode().strip() + raise RuntimeError(f"Interface {interface} not found: {error}") + + # Check if we have necessary permissions + _, stdout, stderr = client.exec_command("which tcpdump") + if stdout.channel.recv_exit_status() != 0: + raise RuntimeError("tcpdump not found on remote system") + + # Build tcpdump command with filter if provided + tcpdump_cmd = f"tcpdump -U -i {interface} -w - " + if capture_filter: + tcpdump_cmd += f"'{capture_filter}'" + + print(f"Starting capture on {ip}:{interface}...") + + # Start tcpdump process on remote machine and get its output + stdin, stdout, stderr = client.exec_command(tcpdump_cmd) + + # Check if tcpdump started successfully (non-blocking check) + time.sleep(1) + if stdout.channel.exit_status_ready(): + error = stderr.read().decode().strip() + raise RuntimeError(f"Failed to start tcpdump: {error}") + + # Create a named pipe (FIFO) for tshark to read from + fifo_path = tempfile.mktemp() + os.mkfifo(fifo_path) + + # Start a process to read from SSH and write to the FIFO + def pipe_ssh_to_fifo(): + try: + with open(fifo_path, 'wb') as fifo: + while True: + data = stdout.read(4096) + if not data: + break + fifo.write(data) + fifo.flush() + except (BrokenPipeError, OSError) as e: + print(f"Error in pipe_ssh_to_fifo: {str(e)}") + finally: + print("Closing FIFO due to error or completion.") + + import threading + thread = threading.Thread(target=pipe_ssh_to_fifo, daemon=True) + thread.start() + + print(f"Capture running. Data available at: {fifo_path}") + print(f"You can now use: tshark -r {fifo_path} -c 100 [options]") + + # Example usage in the context manager + subprocess.run(["tshark", "-r", fifo_path, "-c", "100"]) + + return fifo_path + + except paramiko.AuthenticationException: + raise ConnectionError("Authentication failed. Check username and password.") + except paramiko.SSHException as e: + raise ConnectionError(f"SSH connection error: {str(e)}") + except socket.timeout: + raise ConnectionError(f"Connection timed out after {timeout} seconds") + except Exception as e: + raise RuntimeError(f"Unexpected error: {str(e)}") + + +@function_tool # TODO: not ideal to decorete this context manager. +@contextmanager +def remote_capture_session(ip, username, password, interface, capture_filter="", port=22): + """ + Context manager for remote traffic capture that automatically cleans up resources. + + Usage: + with remote_capture_session("192.168.1.100", "admin", "password", "eth0") as fifo_path: + # Run tshark to read from the FIFO + subprocess.run(["tshark", "-r", fifo_path, "-T", "fields", "-e", "ip.src"]) + """ + fifo_path = None + client = None + + try: + fifo_path = capture_remote_traffic(ip, username, password, interface, + capture_filter=capture_filter, port=port) + yield fifo_path + finally: + if fifo_path and os.path.exists(fifo_path): + try: + os.unlink(fifo_path) + except: + pass + +if __name__ == "__main__": + # Example usage + if len(sys.argv) < 5: + print("Usage: capture_traffic.py [filter]") + sys.exit(1) + + ip = sys.argv[1] + username = sys.argv[2] + password = sys.argv[3] + interface = sys.argv[4] + capture_filter = sys.argv[5] if len(sys.argv) > 5 else "" + + try: + with remote_capture_session(ip, username, password, interface, capture_filter) as fifo_path: + # Keep the script running until interrupted + print("Press Ctrl+C to stop the capture") + while True: + time.sleep(1) + except KeyboardInterrupt: + print("\nCapture stopped") + except Exception as e: + print(f"Error: {str(e)}") + sys.exit(1) \ No newline at end of file