Push

.gitignore

@@ -145,3 +145,4 @@ pentestperf
 cai.png
 cai.dot
 cai*.png
+test_20*

README.md

@@ -361,3 +361,7 @@ CAI is developed by [Alias Robotics](https://aliasrobotics.com) and funded as pa
 - [ ] Tracing
 - [ ] Graph/flow and other abstractions
 - [ ] Plan/router
+
+### ENV Vars
+
+CTF_IN_DOCKER="True"

cai/tools/cli.py

@@ -1,6 +1,11 @@
+"""
+Here are the CLI tools for executing commands, they should usually return
+with run_command, but you can be creative
+"""
 
 from .common import run_command
 
+
 def list_dir(path: str, args: str, ctf=None) -> str:
     """
     List the contents of a directory.
@@ -15,6 +20,7 @@ def list_dir(path: str, args: str, ctf=None) -> str:
     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.
@@ -27,4 +33,4 @@ def cat_file(args: str, file_path: str, ctf=None) -> str:
         str: The output of running the cat command
     """
     command = f'cat {args} {file_path} '
-    return run_command(command, ctf=ctf)
+    return run_command(command, ctf=ctf)

cai/tools/common.py

@@ -1,6 +1,13 @@
-from wasabi import color
-import subprocess
+"""Basic utilities for executing commands inside and outside virtual
+containers. It's important to note that for proper functioning, the
+CTF_IN_DOCKER environment variable must be configured, and commands can
+only be executed through CTF instances with pentestperf within the code.
+See examples/cybersecurity/1_arch_short_picoctf_static_flag.py for an
+example."""
+
 import os
+import subprocess  # nosec B404
+from wasabi import color  # pylint: disable=import-error
 
 
 def _run_ctf(ctf, command, stdout=True):
@@ -13,7 +20,7 @@ def _run_ctf(ctf, command, stdout=True):
         if stdout:
             print("\033[32m" + output + "\033[0m")
         return output
-    except Exception as e:
+    except Exception as e:  # pylint: disable=broad-except
         print(color(f"Error executing CTF command: {e}", fg="red"))
         # exploit_logger.log_error(str(e))
         return f"Error executing CTF command: {str(e)}"
@@ -21,16 +28,18 @@ def _run_ctf(ctf, command, stdout=True):
 
 def _run_attacker_machine(command, stdout=True):
     try:
+        # nosec B602 - shell=True is required for command chaining
         result = subprocess.run(
             command,
-            shell=True,
+            shell=True,  # nosec B602
             capture_output=True,
-            text=True)
+            text=True,
+            check=True)
         output = result.stdout
         if stdout:
             print("\033[32m" + output + "\033[0m")
         return output
-    except Exception as e:
+    except Exception as e:  # pylint: disable=broad-except
         print(color(f"Error executing local command: {e}", fg="red"))
         return f"Error executing local command: {str(e)}"
 
@@ -42,5 +51,4 @@ def run_command(command, ctf=None, stdout=True):
     """
     if os.getenv("CTF_IN_DOCKER", "false").lower() == "true" and ctf:
         return _run_ctf(ctf, command, stdout)
-    else:
-        return _run_attacker_machine(command, stdout)
+    return _run_attacker_machine(command, stdout)

examples/cybersecurity/1_arch_short_picoctf_static_flag.py

@@ -6,8 +6,10 @@ from cai import CAI, Agent
 from cai.tools.cli import list_dir, cat_file
 import pentestperf as ptt
 from wasabi import color
+import os
 
 # 2. Init the CTF environment
+os.environ["CTF_IN_DOCKER"] = "true"
 ctf = ptt.ctf(
     "picoctf_static_flag",
     subnet="192.168.9.0/24",

tests/CAI/test_runs/test_core.py

@@ -1,6 +1,6 @@
 import pytest
 from cai import CAI, Agent
-from tests.mock_client import MockOpenAIClient, create_mock_response
+from .mock_client import MockOpenAIClient, create_mock_response
 from unittest.mock import Mock
 import json
 

tests/test_runs/test_20240402-113647.json

@@ -1 +0,0 @@
-[{"task_id": "02b37e8e-e436-445c-abdc-13e227616e07", "role": "user", "content": "If I have 5 ducks, and lose 2 of them. How many do I have left"}, {"task_id": "02b37e8e-e436-445c-abdc-13e227616e07", "role": "assistant", "content": "Response to user: 3 ducks"}]