david/gensec-westgate-djw2
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

small code fixes

Browse Source
This commit is contained in:
David Westgate 2024-06-15 01:48:12 -07:00
parent 98d287f207
commit 9f75456b7c
3 changed files with 14 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
final/README.md
View File

@ -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*
*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*

5
final/app.py
View File

@ -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)

10
final/tools.py
View File

@ -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()