From 727a64f3e5d2b8cbd7f3d6e8c0d41245816c8753 Mon Sep 17 00:00:00 2001 From: David Westgate Date: Wed, 22 May 2024 16:40:46 -0700 Subject: [PATCH] add tools, update readme, format --- .gitignore | 3 +- hw6/README.md | 14 ++++++- hw6/app.py | 45 ++++++++++++++++------ hw6/tools.py | 104 +++++++++++++++++++++++++++++++------------------- 4 files changed, 112 insertions(+), 54 deletions(-) diff --git a/.gitignore b/.gitignore index ad19a0e..f8156f0 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,5 @@ __pycache__/ rag_data .chromadb *temp* -*downloads* \ No newline at end of file +*downloads* +*wordlist* \ No newline at end of file diff --git a/hw6/README.md b/hw6/README.md index e9ae2ec..58600ce 100644 --- a/hw6/README.md +++ b/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 diff --git a/hw6/app.py b/hw6/app.py index fc83766..ea49d3b 100644 --- a/hw6/app.py +++ b/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("Welcome to the Wi-Fi reconnissance AI assistant. I can perform wifi recon and penetration tasks from the radio on your local machine.") +print( + "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}) - print(result) + result: Dict[str, any] = agent_executor.invoke({"input": line}) + print(result["output"]) else: break diff --git a/hw6/tools.py b/hw6/tools.py index f0279a8..fb5d458 100644 --- a/hw6/tools.py +++ b/hw6/tools.py @@ -11,63 +11,87 @@ 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): -# query: str = Field(description="should be a request to enable monitor mode of the wireless adapter") -# def get_wireless_interface(): -# 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] - +class CrackPassword(BaseModel): + query: str = Field( + description="Should be command line parameters to 'aircrack-ng' to perform some kind of wifi encryption cracking" + ) + +@tool( + "Perform wifi encryption cracking with aircrack-ng", + 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 -# @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 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 + 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") - # def get_wireless_interface(): - # 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] - + query: str = Field( + 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 ,return_direct=False) +@tool( + "Change the state of the wireless adapter mode with airmon-ng", + 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}"]}) - return res + """Return wireless interface information via iwconfig""" + res = shell_tool.run({"commands": [f"iwconfig {params}"]}) + return res