david/gensec-westgate-djw2
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

code cleanup and comments

Browse Source
This commit is contained in:
David Westgate 2024-05-23 20:00:52 -07:00
parent bca136c1c0
commit ca919d9520
2 changed files with 51 additions and 126 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
hw6/app.py
View File

@ -1,55 +1,38 @@
import subprocess
from typing import Dict
import requests
import tempfile
import zipfile
import io
import os
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
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.agents import AgentExecutor, create_react_agent
from langchain.tools import tool
from tools import (
get_wireless_interface,
change_adapter_mode,
# wifi_network_reconnissance,
wifi_encryption_cracking,
# packet_frame_transmission,
# packet_capture_reconnaissance,
deauth_and_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()
"""
Main application logic. We just loop for user input, pull in tools, and create an agent executor with a custom prompt.
"""
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)
tools = []
tools.extend(
[
get_wireless_interface,
change_adapter_mode,
# wifi_network_reconnissance,
wifi_encryption_cracking,
# packet_frame_transmission,
# packet_capture_reconnaissance,
deauth_and_capture,
reconnaissance,
]
@ -73,6 +56,7 @@ 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}")

145
hw6/tools.py
View File

@ -1,27 +1,30 @@
import json
import subprocess
import threading
import requests
import tempfile
import zipfile
import io
import random
import os
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
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_core.pydantic_v1 import BaseModel, Field
from langchain.tools import tool
from langchain import hub
from langchain_community.tools import ShellTool
from dotenv import load_dotenv
from time import time
"""
All of my tools are built to run pre-defined shell tools in particular ways and generally just take CLI parameters as input.
I could have gave my agent 1 simple terminal tool, and a strong prompt mentioning the various kali applications used here and let it figure out the details.
That might have worked with a lot less code, but when done like this, we lessen the chance that a 'sudo' command goes rogue, and we also handle the needed multithreading case in python
"""
shell_tool = ShellTool()
"""
Crack password is a wrapper to run aircrack-ng and expects the LLM to provide the BSSID and packet capture file containing the handshake.
Some extra work could be done here to improve robustness including
1) Handle gracefully if the .cap does not actually contain a handshake
2) Take in a custom wordlist provided by the LLM
OR
3) Dynamically load one or more word list on the fly from the wordlist directory
4) Show some user feedback while cracking - this can take a while =/
"""
class CrackPassword(BaseModel):
json_params: str = Field(
description="""Should be a string of a json object containing
@ -37,12 +40,10 @@ class CrackPassword(BaseModel):
)
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(
{
@ -51,9 +52,11 @@ def wifi_encryption_cracking(json_params: str) -> str:
)
return res
"""Deauth via combination of airodump and airreplay"""
"""
This tool helps us run a combination of airodump-ng and aireplay-ng simultanously to de-auth and capture the handshake.
Because we must run these at the same time, we cannot just re-use the reconnaissance tool; instead we need this tool which runs each process in a seprate thread.
Overall this works ok assuming the client device choosing to re-connect after a de-auth, which you cannot always control.
"""
class DeauthAndCapture(BaseModel):
json_params: str = Field(
@ -69,16 +72,14 @@ class DeauthAndCapture(BaseModel):
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}"
# 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(
{
@ -87,7 +88,6 @@ def deauth_and_capture(json_params: str) -> str:
]
}
)
def run_aireplay(bssid, client, interface):
shell_tool.run(
{
@ -96,45 +96,25 @@ def deauth_and_capture(json_params: str) -> str:
]
}
)
thread1 = threading.Thread(
dumpthread = 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()
playthread = threading.Thread(target=run_aireplay, args=(bssid, client, interface))
dumpthread.start()
playthread.start()
dumpthread.join()
playthread.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
"""
This tools wraps airodump-ng. airodump is tricky to work with - it does not have any parameter to stop on its own so it needs a timing control,
and the SIGKILL is unfortunetly what it takes. It also does not do well with printing to stdout in this enviornment, so we must print to a dump file, and read the contents
of the dump file to get this tool to return what we are interested in. I looked into other approaches like 'iw scan', they had their own problems so we just stick with the aircrack suite
"""
class PacketCapture(BaseModel):
params: str = Field(
description="Should be the interface name on which to perform reconnaissance, followed by any optional command line flag parameters"
)
@tool(
"Perform wifi reconnaissance with airodump-ng. Interface must be in monitor mode before starting",
args_schema=PacketCapture,
@ -143,11 +123,6 @@ class PacketCapture(BaseModel):
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)
# json_obj = json.loads(json_params)
# interface = json_obj["interface"]
# params = json_obj["params"]
# if params == False:
# params = ""
res = shell_tool.run(
{
"commands": [
@ -157,7 +132,6 @@ def reconnaissance(params: str) -> str:
)
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
@ -165,44 +139,13 @@ def reconnaissance(params: str) -> str:
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"
# )
# 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
"""airmon-ng tool"""
"""
Straightfowrad tool wrapping airmon-ng
"""
class ChangeMonitorMode(BaseModel):
params: 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,
@ -211,19 +154,17 @@ class ChangeMonitorMode(BaseModel):
def change_adapter_mode(params: str) -> str:
"""Can pass parameters to airmon-ng to change the mode of the wireless adapter"""
params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM
res = shell_tool.run({"commands": [f"sudo airmon-ng {params}", "sleep 5s"]})
res = shell_tool.run({"commands": [f"sudo airmon-ng {params}", "sleep 5s"]}) # takes time, good idea to sleep
return res
"""iwconfig tool"""
"""
Straightfowrad tool wrapping iwconfig
"""
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"
)
@tool("Get interface information", args_schema=Iwconfig, return_direct=False)
def get_wireless_interface(params: str) -> str:
"""Return wireless interface information via iwconfig"""