diff --git a/final/README.md b/final/README.md index 62a4646..7d3e65d 100644 --- a/final/README.md +++ b/final/README.md @@ -19,8 +19,10 @@ python3 app.py ``` ## Example tests -*Show a summary of network traffic on enp7s0* +*Give a summary of network traffic on my ethernet adapter* -*Show a visual representation of the network traffic on enp7s0* +*Based on ethernet traffic, what applications might I be running* -*Show a list of IP addresses that communicate on enp7s0 and identify the countries of those addresses* \ No newline at end of file +*What countries or regions am I sending traffic to via my ethernet adapter* + +*Show a visual representation of the network traffic on my ethernet adapter* \ No newline at end of file diff --git a/final/app.py b/final/app.py index 42c2e80..6923d85 100644 --- a/final/app.py +++ b/final/app.py @@ -2,7 +2,7 @@ import os from langchain_openai import ChatOpenAI from langchain.agents import AgentExecutor, create_react_agent from langchain.tools import tool -from tools import dalle, tcp_dump, get_wireless_interface, ip_loc, pcap_summary +from tools import dalle, tcp_dump, get_adapter_interface, ip_loc, pcap_summary from langchain import hub from langchain_community.tools import ShellTool from langsmith import Client @@ -22,12 +22,13 @@ client = Client() shell_tool = ShellTool() llm = ChatOpenAI(model_name="gpt-4o", temperature=0) tools = [] -tools.extend([tcp_dump, ip_loc, dalle, pcap_summary, get_wireless_interface]) +tools.extend([tcp_dump, ip_loc, dalle, pcap_summary, get_adapter_interface]) base_prompt = hub.pull("langchain-ai/react-agent-template") prompt = base_prompt.partial( instructions=""" You are a packet analysis assistant. Use any combination of the tools provided to best serve the users request. If the request cannot be served with the tools provided, state why and offer advice on how the user could solve the problem. + Never make assumptions about which network adapter to use, always check the adapter name with the correct tool. """ ) agent = create_react_agent(llm, tools, prompt) diff --git a/final/tools.py b/final/tools.py index c6d0483..38f2a3e 100644 --- a/final/tools.py +++ b/final/tools.py @@ -1,6 +1,7 @@ from datetime import datetime from io import StringIO import sys +from time import sleep from langchain_core.pydantic_v1 import BaseModel, Field from langchain.tools import tool from langchain_community.tools import ShellTool @@ -50,7 +51,7 @@ class Image(BaseModel): @tool( - "Create an image of a network. This tool must take as input a summary of packet information, and not a file name", + "Create an image of a network. This tool must take as input a summary of packet information, and not a file name. If necessary, condense input to 900 characters or less", args_schema=Image, return_direct=False, ) @@ -96,9 +97,9 @@ class Iwconfig(BaseModel): ) -@tool("Get interface information", args_schema=Iwconfig, return_direct=False) -def get_wireless_interface(params: str) -> str: - """Return wireless interface information via iwconfig""" +@tool("Get interface adapter name", args_schema=Iwconfig, return_direct=False) +def get_adapter_interface(params: str) -> str: + """Network interface adapters via iwconfig""" params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM print("params ", params) res = shell_tool.run({"commands": [f"iwconfig {params}"]}) @@ -109,6 +110,7 @@ def get_wireless_interface(params: str) -> str: def ip_loc(address): """Get information from an ip address, including geolocation. Takes as a paramater an ip address. Do not use this tool with IP adresses in a reserve range or on LAN""" url = f"http://ipwho.is/{address}" + sleep(2) response = requests.get(url) if response.status_code == 200: return response.json()