From 4bce6cfc067b4d1dc8c98d743fe7844ea157ef6e Mon Sep 17 00:00:00 2001 From: luijait2 Date: Thu, 9 Jan 2025 17:14:06 +0000 Subject: [PATCH] Push --- .gitignore | 1 + README.md | 4 ++++ cai/tools/cli.py | 8 ++++++- cai/tools/common.py | 24 ++++++++++++------- .../1_arch_short_picoctf_static_flag.py | 2 ++ tests/{ => CAI/test_runs}/__init__.py | 0 tests/{ => CAI/test_runs}/mock_client.py | 0 tests/{ => CAI/test_runs}/test_core.py | 2 +- tests/{ => CAI/test_runs}/test_util.py | 0 tests/test_runs/test_20240402-113647.json | 1 - 10 files changed, 31 insertions(+), 11 deletions(-) rename tests/{ => CAI/test_runs}/__init__.py (100%) rename tests/{ => CAI/test_runs}/mock_client.py (100%) rename tests/{ => CAI/test_runs}/test_core.py (98%) rename tests/{ => CAI/test_runs}/test_util.py (100%) delete mode 100644 tests/test_runs/test_20240402-113647.json diff --git a/.gitignore b/.gitignore index 8a225e24..74dc99a6 100644 --- a/.gitignore +++ b/.gitignore @@ -145,3 +145,4 @@ pentestperf cai.png cai.dot cai*.png +test_20* diff --git a/README.md b/README.md index 49474e4e..fb7e1525 100644 --- a/README.md +++ b/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" diff --git a/cai/tools/cli.py b/cai/tools/cli.py index 44ad72c5..38deaa2c 100644 --- a/cai/tools/cli.py +++ b/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) \ No newline at end of file + return run_command(command, ctf=ctf) diff --git a/cai/tools/common.py b/cai/tools/common.py index d8baff8f..5a5f01e9 100644 --- a/cai/tools/common.py +++ b/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) diff --git a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py b/examples/cybersecurity/1_arch_short_picoctf_static_flag.py index 241de965..6bc6d32e 100644 --- a/examples/cybersecurity/1_arch_short_picoctf_static_flag.py +++ b/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", diff --git a/tests/__init__.py b/tests/CAI/test_runs/__init__.py similarity index 100% rename from tests/__init__.py rename to tests/CAI/test_runs/__init__.py diff --git a/tests/mock_client.py b/tests/CAI/test_runs/mock_client.py similarity index 100% rename from tests/mock_client.py rename to tests/CAI/test_runs/mock_client.py diff --git a/tests/test_core.py b/tests/CAI/test_runs/test_core.py similarity index 98% rename from tests/test_core.py rename to tests/CAI/test_runs/test_core.py index b2777bdb..6c6cb4b4 100644 --- a/tests/test_core.py +++ b/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 diff --git a/tests/test_util.py b/tests/CAI/test_runs/test_util.py similarity index 100% rename from tests/test_util.py rename to tests/CAI/test_runs/test_util.py diff --git a/tests/test_runs/test_20240402-113647.json b/tests/test_runs/test_20240402-113647.json deleted file mode 100644 index b020a09d..00000000 --- a/tests/test_runs/test_20240402-113647.json +++ /dev/null @@ -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"}]