From cc06d9d5b202563ac98f988a1162fbb239958f52 Mon Sep 17 00:00:00 2001
From: luijait2 <luijait@aliasrobotics.com>
Date: Thu, 9 Jan 2025 16:59:43 +0000
Subject: [PATCH] Added CTF_IN_DOCKER variable

---
 cai/core.py                                   |  7 +--
 cai/tools/common.py                           | 17 +++++--
 .../1_arch_short_picoctf_static_flag.py       | 46 +++++++++++++++++++
 3 files changed, 63 insertions(+), 7 deletions(-)
 create mode 100644 examples/cybersecurity/1_arch_short_picoctf_static_flag.py

diff --git a/cai/core.py b/cai/core.py
index 35a4b597..93cbddc4 100644
--- a/cai/core.py
+++ b/cai/core.py
@@ -14,7 +14,7 @@ import copy
 import json
 from collections import defaultdict
 from typing import List
-
+import os
 # Package/library imports
 from openai import OpenAI  # pylint: disable=import-error
 
@@ -31,7 +31,6 @@ from .types import (
     Result,
 )
 
-from .tools.common import run_command
 __CTX_VARS_NAME__ = "context_variables"
 
 
@@ -150,9 +149,11 @@ class CAI:
             func = function_map[name]
             # pass context_variables to agent functions
             if __CTX_VARS_NAME__ in func.__code__.co_varnames:
-                args[__CTX_VARS_NAME__] = context_variables 
+                args[__CTX_VARS_NAME__] = context_variables
             if self.ctf:
                 args["ctf"] = self.ctf
+            else:
+                os.environ["CTF_IN_DOCKER"] = "false"
             raw_result = function_map[name](**args)
 
             result: Result = self.handle_function_result(raw_result, debug)
diff --git a/cai/tools/common.py b/cai/tools/common.py
index e8e4a78a..d8baff8f 100644
--- a/cai/tools/common.py
+++ b/cai/tools/common.py
@@ -2,9 +2,11 @@ from wasabi import color
 import subprocess
 import os
 
+
 def _run_ctf(ctf, command, stdout=True):
     try:
-        # Ensure the command is executed in a shell that supports command chaining
+        # Ensure the command is executed in a shell that supports command
+        # chaining
         output = ctf.get_shell(command)
         # exploit_logger.log_ok()
 
@@ -16,9 +18,14 @@ def _run_ctf(ctf, command, stdout=True):
         # exploit_logger.log_error(str(e))
         return f"Error executing CTF command: {str(e)}"
 
+
 def _run_attacker_machine(command, stdout=True):
     try:
-        result = subprocess.run(command, shell=True, capture_output=True, text=True)
+        result = subprocess.run(
+            command,
+            shell=True,
+            capture_output=True,
+            text=True)
         output = result.stdout
         if stdout:
             print("\033[32m" + output + "\033[0m")
@@ -27,11 +34,13 @@ def _run_attacker_machine(command, stdout=True):
         print(color(f"Error executing local command: {e}", fg="red"))
         return f"Error executing local command: {str(e)}"
 
+
 def run_command(command, ctf=None, stdout=True):
     """
-    Run command either in CTF container or on local attacker machine based on CTF_IN_DOCKER flag
+    Run command either in CTF container or
+    on local attacker machine based on CTF_IN_DOCKER flag
     """
-    if os.getenv("CTF_IN_DOCKER").lower() == "true" and ctf:
+    if os.getenv("CTF_IN_DOCKER", "false").lower() == "true" and ctf:
         return _run_ctf(ctf, command, stdout)
     else:
         return _run_attacker_machine(command, stdout)
diff --git a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py
new file mode 100644
index 00000000..241de965
--- /dev/null
+++ b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py
@@ -0,0 +1,46 @@
+"""
+This is a short example of how we can solve ctf scenarios programmatically. With CAI framework
+"""
+# 1. Instance CAI, pentestperf and tools
+from cai import CAI, Agent
+from cai.tools.cli import list_dir, cat_file
+import pentestperf as ptt
+from wasabi import color
+
+# 2. Init the CTF environment
+ctf = ptt.ctf(
+    "picoctf_static_flag",
+    subnet="192.168.9.0/24",
+    container_name="ancient",
+    ip_address="192.168.9.52"
+)
+ctf.start_ctf()
+challenges = list(ctf.get_challenges().keys())
+challenge = challenges[0]
+
+# 3. Instance the cybersecurity AI Framework client
+client = CAI(ctf=ctf)
+
+# 4. Instance the agent/agents to solve the CTF
+ctf_agent = Agent(
+    model="qwen2.5:14b",
+    name="A file listing tool",
+    instructions="You are an agent tailored to list and read files within a file system.",
+    functions=[list_dir, cat_file],  # This code can execute by LLM
+    # parallel_tool_calls=True, Enable multi function calling
+)
+
+
+messages = [{
+    "role": "user",
+    "content": "Instructions: " + ctf.get_instructions() +
+               "\nChallenge: " + ctf.get_challenges()[challenge] +
+               "\nTechniques: " + ctf.get_techniques() +
+               "\nExtract and return only the flag"
+}]
+
+# 5. Run the swarm
+response = client.run(agent=ctf_agent, messages=messages, debug=True)
+print(response.messages[-1]["content"])
+
+ctf.stop_ctf()