This repository has been archived on 2025-04-28. You can view files and clone it, but cannot push or open issues or pull requests.
gensec-westgate-djw2/hw6/app.py
David Westgate f7eae67744 added tools
2024-05-22 15:59:39 -07:00

55 lines
2.0 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 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