add tools, update readme, format

.gitignore

@@ -7,3 +7,4 @@ rag_data
 .chromadb
 *temp*
 *downloads*
+*wordlist*

hw6/README.md

@@ -1,6 +1,13 @@
 ###### David Westgate 24 May 2024
 ## HW6 for gensec
+This LLM agent application is a tool suite which intends to help with reconnissance and penetration testing of wifi networks using the radio on the local machine.
 
+## Prerequisites
+This application assumes you are running a real linux system, with a wifi radio properly installed and working. It also assumes you have the [aircrack-ng toolsuite](https://www.aircrack-ng.org/) installed (common on kali systems).
+
+The [Cracking Wifi guide](https://dkmcgrath.github.io/courses/netsec/crack_wifi.html) provided by Professor Kevin McGrath serves as an inpiration for this application.
+
+Note: This application is for penetration testing only on network for which you are authorized to perform suchs tests. Use of this application on any other networks may be illegal. 
 
 ### Setup + Run
 Install python3, then
@@ -10,5 +17,10 @@ pip install -r requirements.txt
 cp .env.example .env #fill in env file with key
 python3 app.py
 ```
+### Optional: Download word lists for cracking
+```
+mkdir wordlist
+curl -o wordlist/rockyou.txt https://github.com/brannondorsey/naive-hashcat/releases/download/data/rockyou.txt
+```
 
-### Results
+## Results

hw6/app.py

@@ -1,4 +1,5 @@
 import subprocess
+from typing import Dict
 import requests
 import tempfile
 import zipfile
@@ -11,7 +12,14 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain.agents import AgentExecutor, create_react_agent, load_tools
 from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
 from langchain.tools import tool
-from tools import get_wireless_interface, change_adapter_mode
+from tools import (
+    get_wireless_interface,
+    change_adapter_mode,
+    wifi_encryption_cracking,
+    packet_capture_reconnaissance,
+    wifi_encryption_cracking,
+    packet_frame_transmission,
+)
 
 from langchain import hub
 from langchain_community.tools import ShellTool
@@ -25,30 +33,43 @@ shell_tool = ShellTool()
 
 llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
 tools = []
-tools.extend([change_adapter_mode, get_wireless_interface])
+tools.extend(
+    [
+        get_wireless_interface,
+        change_adapter_mode,
+        wifi_encryption_cracking,
+        packet_capture_reconnaissance,
+        wifi_encryption_cracking,
+        packet_frame_transmission,
+    ]
+)
 base_prompt = hub.pull("langchain-ai/react-agent-template")
-prompt = base_prompt.partial(instructions=
+prompt = base_prompt.partial(
-                             """
+    instructions="""
                                 You are a wireless network penetration testing assistant
                                 Answer the user's request by utilizing the available tools including iwconfig, airmon-ng, airodump-ng and aircrack-ng. 
                                 As necessary, combine the use of various tools to fufill the request
                                 If a tool is not availble to answer the users request, reject the request and provide the reason.
                              
-                             """)
+                             """
-agent = create_react_agent(llm,tools,prompt)
+)
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
+agent = create_react_agent(llm, tools, prompt)
+agent_executor = AgentExecutor(
+    agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
+)
 
 
-
+print(
-print("Welcome to the Wi-Fi reconnissance AI assistant. I can perform wifi recon and penetration tasks from the radio on your local machine.")
+    "Welcome to the Wi-Fi reconnissance AI assistant. I can perform wifi recon and penetration tasks from the radio on your local machine."
+)
 print(f"I am configured with these tools")
 for tool in tools:
-  print(f'  Tool: {tool.name} = {tool.description}')
+    print(f"  Tool: {tool.name} = {tool.description}")
 
 while True:
     line = input("llm>> ")
     if line:
-        result = agent_executor.invoke({"input":line})
+        result: Dict[str, any] = agent_executor.invoke({"input": line})
-        print(result)
+        print(result["output"])
     else:
         break

hw6/tools.py

@@ -11,62 +11,86 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain.agents import AgentExecutor, create_react_agent, load_tools
 from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
 from langchain.tools import tool
-
-
 from langchain import hub
 from langchain_community.tools import ShellTool
-
-
 from dotenv import load_dotenv
 from time import time
+
 shell_tool = ShellTool()
 
 
-# class EnableMonitorMode(BaseModel):
+class CrackPassword(BaseModel):
-#     query: str = Field(description="should be a request to enable monitor mode of the wireless adapter")
+    query: str = Field(
-#     def get_wireless_interface():
+        description="Should be command line parameters to 'aircrack-ng' to perform some kind of wifi encryption cracking"
-#         result = subprocess.run(['iw', 'dev'], stdout=subprocess.PIPE)
+    )
-#         output = result.stdout.decode('utf-8')
+
-#         for line in output.splitlines():
+@tool(
-#             if 'Interface' in line:
+    "Perform wifi encryption cracking with aircrack-ng",
-#                 return line.split()[1]
+    args_schema=CrackPassword,
+    return_direct=False,
+)
+def wifi_encryption_cracking(params: str) -> str:
+    """Can pass parameters to aircrack-ng to perform wifi encryption cracking"""
+    res = shell_tool.run({"commands": [f"aircrack-ng {params}"]})
+    return res
 
 
+class PacketTransmission(BaseModel):
+    query: 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"""
+    res = shell_tool.run({"commands": [f"sudo aireplay-ng {params}"]})
+    return res
+
+
+class PacketCapture(BaseModel):
+    query: 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}"]})
+    return res
 
-# @tool("enable monitor mode", args_schema=EnableMonitorMode ,return_direct=False)
-# def enable_monitor_mode(interface_name: str) -> str:
-#     """Can enable monitor mode"""
-#     if interface_name.find('mon'):
-#          res = "Montior mode already enabled"
-#     else:
-#         res = shell_tool.run({"commands": ["sudo airmon-ng check kill", f"sudo airmon-ng start {interface_name}"]})
-#     return res
 
 class ChangeMonitorMode(BaseModel):
-    query: str = Field(description="Should be command line parameters to 'airmon-ng' to change the state of a given wireless iterface mode")
+    query: str = Field(
-    # def get_wireless_interface():
+        description="Should be command line parameters to 'airmon-ng' to change the state of a given wireless iterface mode"
-    #     result = subprocess.run(['iw', 'dev'], stdout=subprocess.PIPE)
+    )
-    #     output = result.stdout.decode('utf-8')
-    #     for line in output.splitlines():
-    #         if 'Interface' in line:
-    #             return line.split()[1]
 
-
+@tool(
-
+    "Change the state of the wireless adapter mode with airmon-ng",
-@tool("Change the state of the wireless adapter mode with airmon-ng", args_schema=ChangeMonitorMode ,return_direct=False)
+    args_schema=ChangeMonitorMode,
+    return_direct=False,
+)
 def change_adapter_mode(params: str) -> str:
     """Can pass parameters to airmon-ng to change the mode of the wireless adapter"""
-    # if interface_name.find('mon'):
-    #      res = "Montior mode already enabled"
-    # else:
     res = shell_tool.run({"commands": [f"sudo airmon-ng {params}"]})
     return res
 
+
 class Iwconfig(BaseModel):
-    params: str = Field(description="should be command line parameters to 'iwconfig', if needed")
+    params: str = Field(
+        description="should be command line parameters to 'iwconfig', if needed"
+    )
 
-
+@tool("Get interface information", args_schema=Iwconfig, return_direct=False)
-@tool("Get interface information", args_schema=Iwconfig ,return_direct=False)
 def get_wireless_interface(params: str) -> str:
     """Return wireless interface information via iwconfig"""
     res = shell_tool.run({"commands": [f"iwconfig {params}"]})