add tools, update readme, format

This commit is contained in:
David Westgate 2024-05-22 16:40:46 -07:00
parent f7eae67744
commit 727a64f3e5
4 changed files with 112 additions and 54 deletions

1
.gitignore vendored
View File

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

View File

@ -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

View File

@ -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

View File

@ -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
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")
# 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