Add CTF_IN_DOCKER

.env.example

@@ -0,0 +1,2 @@
+OPENAI_API_KEY=
+CTF_IN_DOCKER=

README.md

@@ -355,7 +355,8 @@ CAI is developed by [Alias Robotics](https://aliasrobotics.com) and funded as pa
 # Plan of development
 - [x] Dev container
 - [x] pre-commit hooks
-- [ ] A first example with agents - picoctf_static_flag
+- [x] A first example with agents - picoctf_static_flag
+  - [ ] ...
 - [ ] CI
 - [ ] Tracing
 - [ ] Graph/flow and other abstractions

cai/core.py

@@ -31,6 +31,7 @@ from .types import (
     Result,
 )
 
+from .tools.common import run_command
 __CTX_VARS_NAME__ = "context_variables"
 
 
@@ -42,10 +43,12 @@ class CAI:
     def __init__(self,
                  client=None,
                  base_url="http://host.docker.internal:8000/v1",
-                 api_key="alias"):
+                 api_key="alias",
+                 ctf=None):
         if not client:
             client = OpenAI(base_url=base_url, api_key=api_key)
         self.client = client
+        self.ctf = ctf
 
     def get_chat_completion(  # pylint: disable=too-many-arguments
         self,
@@ -147,7 +150,9 @@ 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
             raw_result = function_map[name](**args)
 
             result: Result = self.handle_function_result(raw_result, debug)

cai/tools/cli.py

@@ -0,0 +1,30 @@
+
+from .common import run_command
+
+def list_dir(path: str, args: str, ctf=None) -> str:
+    """
+    List the contents of a directory.
+    by def .
+    Args:
+        path: The directory path to list contents from
+        args: Additional arguments to pass to the ls command
+
+    Returns:
+        str: The output of running the ls command
+    """
+    command = f'ls {path} {args}'
+    return run_command(command, ctf=ctf)
+
+def cat_file(args: str, file_path: str, ctf=None) -> str:
+    """
+    Display the contents of a file.
+
+    Args:
+        args: Additional arguments to pass to the cat command
+        file_path: Path to the file to display contents of
+
+    Returns:
+        str: The output of running the cat command
+    """
+    command = f'cat {args} {file_path} '
+    return run_command(command, ctf=ctf)

cai/tools/common.py

@@ -0,0 +1,37 @@
+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
+        output = ctf.get_shell(command)
+        # exploit_logger.log_ok()
+
+        if stdout:
+            print("\033[32m" + output + "\033[0m")
+        return output
+    except Exception as e:
+        print(color(f"Error executing CTF command: {e}", fg="red"))
+        # 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)
+        output = result.stdout
+        if stdout:
+            print("\033[32m" + output + "\033[0m")
+        return output
+    except Exception as e:
+        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
+    """
+    if os.getenv("CTF_IN_DOCKER").lower() == "true" and ctf:
+        return _run_ctf(ctf, command, stdout)
+    else:
+        return _run_attacker_machine(command, stdout)

examples/cybersecurity/1_picoctf_static_flag.py

@@ -4,11 +4,9 @@ from wasabi import color
 
 client = CAI()
 
-
 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()
 
@@ -20,7 +18,6 @@ def run_ctf(ctf, command, stdout=True):
         # exploit_logger.log_error(str(e))
         return f"Error executing CTF command: {str(e)}"
 
-
 def list_dir(path: str, args: str) -> str:
     """
     List the contents of a directory.
@@ -36,7 +33,6 @@ def list_dir(path: str, args: str) -> str:
     command = f'ls {path} {args}'
     return run_ctf(ctf, command)
 
-
 def cat_file(args: str, file_path: str) -> str:
     """
     Display the contents of a file.
@@ -61,7 +57,6 @@ def cat_file(args: str, file_path: str) -> str:
 #      else:
 #         return "Flag Not Found: " + flag
 
-
 listing_agent = Agent(
     model="qwen2.5:14b",
     name="A file listing tool",
@@ -84,11 +79,11 @@ listing_agent = Agent(
 
 # PentestPerf setup for "picoctf_static_flag"
 ctf = ptt.ctf(
-    "picoctf_static_flag",
-    subnet="192.168.9.0/24",
-    container_name="ancient",
-    ip_address="192.168.9.52",
-)
+        "picoctf_static_flag",
+        subnet="192.168.9.0/24",
+        container_name="ancient",
+        ip_address="192.168.9.52",
+    )
 ctf.start_ctf()
 flag = ctf.get_flag()
 
@@ -100,14 +95,14 @@ listing_agent.functions.append(cat_file)
 challenges = list(ctf.get_challenges().keys())
 challenge = challenges[0]
 
-messages = [{"role": "user", "content": "Instructions: " + ctf.get_instructions()
-             + "\nChallenge: " + ctf.get_challenges()[challenge]
-             + "\nTechniques: " + ctf.get_techniques()
-             + "\nExtract and return only the flag"}]
+messages = [{"role": "user", "content": "Instructions: " + ctf.get_instructions() 
+    + "\nChallenge: " + ctf.get_challenges()[challenge]
+    + "\nTechniques: " + ctf.get_techniques()
+    + "\nExtract and return only the flag"}]
 response = client.run(agent=listing_agent, messages=messages, debug=True)
 print(response.messages[-1]["content"])
 
-# if ctf.check_flag(response.messages[-1]["content"], ctf.get_challenges()[challenge]):
-#    print(color(f"Flag found: {response.messages[-1]['content']}", fg="green"))
+#if ctf.check_flag(response.messages[-1]["content"], ctf.get_challenges()[challenge]):
+#    print(color(f"Flag found: {response.messages[-1]['content']}", fg="green"))    
 
 ctf.stop_ctf()