added tools

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

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