diff --git a/hw6/app.py b/hw6/app.py index e69de29..fc83766 100644 --- a/hw6/app.py +++ b/hw6/app.py @@ -0,0 +1,54 @@ +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 tools import get_wireless_interface, change_adapter_mode + +from langchain import hub +from langchain_community.tools import ShellTool + + +from dotenv import load_dotenv +from time import time + +load_dotenv() +shell_tool = ShellTool() + +llm = ChatOpenAI(model_name="gpt-4o", temperature=0) +tools = [] +tools.extend([change_adapter_mode, get_wireless_interface]) +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 + 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) + + + +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}') + +while True: + line = input("llm>> ") + if line: + result = agent_executor.invoke({"input":line}) + print(result) + else: + break diff --git a/hw6/tools.py b/hw6/tools.py new file mode 100644 index 0000000..f0279a8 --- /dev/null +++ b/hw6/tools.py @@ -0,0 +1,73 @@ +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 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] + + + +# @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] + + + +@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") + + +@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