some improvement to tools

This commit is contained in:
David Westgate 2024-05-23 12:00:52 -07:00
parent e6ef898467
commit 9c7dece545
4 changed files with 78 additions and 28 deletions

3
.gitignore vendored
View File

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

View File

@ -15,13 +15,19 @@ Install python3, then
cd hw6
pip install -r requirements.txt
cp .env.example .env #fill in env file with key
mkdir dumps
mkdir wordlist
python3 app.py
```
### Optional: Download word lists for cracking
### Download word lists for cracking
```
mkdir wordlist
curl -o wordlist/rockyou.txt https://github.com/brannondorsey/naive-hashcat/releases/download/data/rockyou.txt
```
## Test
*Find the password of the GenSec wifi network*
For this test:
* Ensure you have a wifi network running and reachable from your computer with the name NetSec
* Ensure it has a password found in the wordlist dictionary
* Ensure a client device is connected
*Find the password of the NetSec wifi network*

View File

@ -18,16 +18,24 @@ from tools import (
wifi_network_reconnissance,
wifi_encryption_cracking,
packet_frame_transmission,
packet_capture_reconnaissance
)
from langchain import hub
from langchain_community.tools import ShellTool
from langsmith import Client
from dotenv import load_dotenv
from time import time
load_dotenv()
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = f"LangSmith Introduction"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
client = Client()
shell_tool = ShellTool()
llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
@ -36,17 +44,18 @@ tools.extend(
[
get_wireless_interface,
change_adapter_mode,
wifi_network_reconnissance,
#wifi_network_reconnissance,
wifi_encryption_cracking,
packet_frame_transmission,
packet_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 tools including iwconfig, airmon-ng, airodump-ng and aircrack-ng.
As necessary, combine the use of various tools to fufill the request
Answer the user's request by utilizing the available any combinations of tools available including iwconfig, airmon-ng, airodump-ng, and aircrack-ng.
If a tool is not availble to answer the users request, reject the request and provide the reason.
"""

View File

@ -3,6 +3,7 @@ import requests
import tempfile
import zipfile
import io
import random
import os
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
@ -18,10 +19,12 @@ from time import time
shell_tool = ShellTool()
class CrackPassword(BaseModel):
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
@ -50,6 +53,7 @@ class CrackPassword(BaseModel):
)
def wifi_encryption_cracking(params: str) -> str:
"""Can pass parameters to aircrack-ng to perform wifi encryption cracking"""
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
res = shell_tool.run({"commands": [f"aircrack-ng {params}"]})
return res
@ -59,6 +63,7 @@ class PacketTransmission(BaseModel):
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,
@ -66,42 +71,67 @@ class PacketTransmission(BaseModel):
)
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 or packet capture"
# )
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.
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",
args_schema=PacketCapture,
return_direct=False,
)
def packet_capture_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"""
hash = random.getrandbits(16)
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
res = shell_tool.run(
{
"commands": [
f"sudo timeout -s SIGKILL 15 airodump-ng --output-format csv --write dumps/dump-{hash} {params} "
]
}
)
fp = open(f"dumps/dump-{hash}-01.csv", "r")
contents = fp.read()
fp.close()
return contents
# @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 timeout 15s airodump-ng {params}"]})
# return res
class IwScan(BaseModel):
interface: str = Field(
description="Should be a wireless interface name, used as a paramater to 'iw' to scan for wifi networks"
)
network: str = Field(
description="Should be the name or SSID of the wifi network you are interested in"
)
# network: str = Field(
# description="Should be the name or SSID of the wifi network you are interested in"
# )
@tool(
"Perform wifi scanning with iw",
"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,
return_direct=False,
)
def wifi_network_reconnissance(interface: str, network: 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"""
res = shell_tool.run({"commands": [f'sudo iw {interface} scan | grep -B 9 -A 206 "{network}"']})
interface = interface.replace("`", "").replace("\n", "") # fix buggy input from LLM
res = shell_tool.run(
{
"commands": [
f"sudo ifconfig {interface} up",
"sleep 5",
f'sudo iw {interface} scan | grep -B 9 -A 206 "NetSec"',
]
}
) # TODO: fix error when passing network param
return res
@ -118,17 +148,21 @@ class ChangeMonitorMode(BaseModel):
)
def change_adapter_mode(params: str) -> str:
"""Can pass parameters to airmon-ng to change the mode of the wireless adapter"""
res = shell_tool.run({"commands": [f"sudo airmon-ng {params}"]})
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
res = shell_tool.run({"commands": [f"sudo airmon-ng {params}", "sleep 5s"]})
return res
class Iwconfig(BaseModel):
params: str = Field(
description="should be command line parameters to 'iwconfig'. If none are needed, this should be left blank"
description="should be command line parameters to 'iwconfig'. If none are needed, this should be left as an empty string"
)
@tool("Get interface information", args_schema=Iwconfig, return_direct=False)
def get_wireless_interface(params: str) -> str:
"""Return wireless interface information via iwconfig"""
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
print("params ", params)
res = shell_tool.run({"commands": [f"iwconfig {params}"]})
return res