116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
import subprocess
|
|
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.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 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
|
|
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,
|
|
)
|
|
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):
|
|
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"""
|
|
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"
|
|
)
|
|
|
|
@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 -s SIGINT 15s airodump-ng {params}"]})
|
|
return res
|
|
|
|
|
|
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,
|
|
return_direct=False,
|
|
)
|
|
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}"]})
|
|
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"
|
|
)
|
|
|
|
@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
|