fix issue with pcap_summary

This commit is contained in:
David Westgate 2024-06-14 13:17:09 -07:00
parent fa5a30db6e
commit 8416de3273

View File

@ -1,4 +1,6 @@
from datetime import datetime from datetime import datetime
from io import StringIO
import sys
from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.tools import tool from langchain.tools import tool
from langchain_community.tools import ShellTool from langchain_community.tools import ShellTool
@ -76,7 +78,15 @@ class Summary(BaseModel):
def pcap_summary(file: str) -> str: def pcap_summary(file: str) -> str:
"""Must pass the path to a pcap file to return a summary of packet information and traffic data""" """Must pass the path to a pcap file to return a summary of packet information and traffic data"""
cap: PacketList = rdpcap(file) cap: PacketList = rdpcap(file)
return cap.summary()
# Workaround to get `summary` to a variable. It normally prints to stdout and returns nothing
stdout_capture = StringIO()
save_stdout = sys.stdout
sys.stdout = stdout_capture
cap.summary()
sys.stdout = save_stdout
return stdout_capture.getvalue()
# From hw6 # From hw6
@ -97,7 +107,7 @@ def get_wireless_interface(params: str) -> str:
@tool @tool
def ip_loc(address): def ip_loc(address):
"""Get information from an ip address, including geolocation. Takes as a paramater an ip 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}" url = f"http://ipwho.is/{address}"
response = requests.get(url) response = requests.get(url)
if response.status_code == 200: if response.status_code == 200: