add wordlist file loading to cracking

hw6/README.md

@@ -23,4 +23,5 @@ mkdir wordlist
 curl -o wordlist/rockyou.txt https://github.com/brannondorsey/naive-hashcat/releases/download/data/rockyou.txt
 ```
 
-## Results
+## Test
+*Find the password of the GenSec wifi network*

hw6/tools.py

@@ -18,11 +18,30 @@ from time import time
 
 shell_tool = ShellTool()
 
-
 class CrackPassword(BaseModel):
-    query: str = Field(
+    params: str = Field(
         description="Should be command line parameters to 'aircrack-ng' to perform some kind of wifi encryption cracking"
     )
+    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",
@@ -36,7 +55,7 @@ def wifi_encryption_cracking(params: str) -> str:
 
 
 class PacketTransmission(BaseModel):
-    query: str = Field(
+    params: str = Field(
         description="Should be command line parameters to 'aireplay-ng' to perform some kind of wifi frame or packet transmission"
     )
 
@@ -52,28 +71,27 @@ def packet_frame_transmission(params: str) -> str:
 
 
 class PacketCapture(BaseModel):
-    query: str = Field(
+    params: str = Field(
         description="Should be command line parameters to 'airodump-ng' to perform some kind of wifi reconnaissance or packet capture"
     )
 
-
 @tool(
     "Perform packet capture or wifi reconnaissance with airodump-ng",
     args_schema=PacketCapture,
     return_direct=False,
 )
-
 def packet_capture_reconnaissance(params: str) -> str:
     """Can pass parameters to airodump-ng to perform packet capture or wifi reconnaissance"""
-    res = shell_tool.run({"commands": [f"sudo airodump-ng {params}"]})
+    res = shell_tool.run({"commands": [f"sudo timeout -s SIGINT 15s airodump-ng {params}"]})
     return res
 
 
 class ChangeMonitorMode(BaseModel):
-    query: str = Field(
+    params: str = Field(
-        description="Should be command line parameters to 'airmon-ng' to change the state of a given wireless iterface mode"
+        description="Should be command line parameters to 'airmon-ng' to change the state of a given wireless iterface mode."
     )
 
+
 @tool(
     "Change the state of the wireless adapter mode with airmon-ng",
     args_schema=ChangeMonitorMode,
@@ -87,7 +105,7 @@ def change_adapter_mode(params: str) -> str:
 
 class Iwconfig(BaseModel):
     params: str = Field(
-        description="should be command line parameters to 'iwconfig', if needed"
+        description="should be command line parameters to 'iwconfig'. If none are needed, this should be left blank"
     )
 
 @tool("Get interface information", args_schema=Iwconfig, return_direct=False)