Add a first example CTF

Peer coded with @Mery-Sanz, @Lidia9, @vmayoral

Signed-off-by: luijait2 <luijait@aliasrobotics.com>

cai/core.py

@@ -39,8 +39,10 @@ class CAI:
     Main class for the CAI library.
     """
 
-    def __init__(self, client=None,
-                 base_url="http://localhost:8000/v1", api_key="alias"):
+    def __init__(self,
+                 client=None,
+                 base_url="http://host.docker.internal:8000/v1",
+                 api_key="alias"):
         if not client:
             client = OpenAI(base_url=base_url, api_key=api_key)
         self.client = client

examples/cybersecurity/1_picoctf_static_flag.py

@@ -1,42 +1,113 @@
 from cai import CAI, Agent
 import pentestperf as ptt
+from wasabi import color
 
 client = CAI()
 
-english_agent = Agent(
+
+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 list_dir(path: str, args: str) -> str:
+    """
+    List the contents of a directory.
+
+    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
+    """
+    global ctf
+    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.
+
+    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
+    """
+    global ctf
+    command = f'cat {args} {file_path} '
+    return run_ctf(ctf, command)
+
+# def flag_checker(flag: str) -> str:
+#      """
+#      This tool is used to check if the guessed flag is valid.
+#      """
+#      if ctf.get_flag(ctf.get_challenges()[challenge]) == flag:
+#         return "Flag Found: " + flag
+#      else:
+#         return "Flag Not Found: " + flag
+
+
+listing_agent = Agent(
     model="qwen2.5:14b",
     name="A file listing tool",
-    instructions="You only speak English. If another language is detected, invoke transfer_to_spanish_agent.",
-    # instructions="You only speak English.",
-    # tool_choice="required",  # not working with ollama and qwen2.5
+    instructions="You are an agent tailored to list and read files within a file system.",
+    # parallel_tool_calls=True,
 )
 
-spanish_agent = Agent(
-    name="Spanish Agent",
-    instructions="You only speak Spanish.",
+# flag_extractor = Agent(
+#     model="qwen2.5:14b",
+#     name="A flag extractor",
+#     instructions="You are an agent tailored to extract only the flag from a given content.",
+# )
+
+# def transfer_to_flag_extractor_agent():
+#     """Transfer to flag extractor."""
+#     return flag_extractor
+
+
+#######################################
+
+# 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",
 )
+ctf.start_ctf()
+flag = ctf.get_flag()
 
 
-def transfer_to_spanish_agent():
-    """Transfer spanish speaking users immediately."""
-    return spanish_agent
+listing_agent.functions.append(list_dir)
+listing_agent.functions.append(cat_file)
+# listing_agent.functions.append(flag_checker)
 
+challenges = list(ctf.get_challenges().keys())
+challenge = challenges[0]
 
-english_agent.functions.append(transfer_to_spanish_agent)
-messages = [{"role": "user", "content": "Hola, cómo estás?"}]
-
-
-response = client.run(agent=english_agent, messages=messages, debug=True)
+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"))
 
-# # 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",
-#     )
-# ctf.start_ctf()
-# flag = ctf.get_flag()
-# ctf.stop_ctf()
+ctf.stop_ctf()

setup.cfg

@@ -19,6 +19,7 @@ install_requires =
     tqdm
     pre-commit
     instructor
+    wasabi
     pentestperf @ git+https://gitlab:gldt-1NfZ2XRWdHNip4K45auL@gitlab.com/aliasrobotics/alias_research/pentestperf@main
 python_requires = >=3.10