diff --git a/final/tools.py b/final/tools.py index e9cdeaf..c6d0483 100644 --- a/final/tools.py +++ b/final/tools.py @@ -1,4 +1,6 @@ from datetime import datetime +from io import StringIO +import sys from langchain_core.pydantic_v1 import BaseModel, Field from langchain.tools import tool from langchain_community.tools import ShellTool @@ -76,7 +78,15 @@ class Summary(BaseModel): def pcap_summary(file: str) -> str: """Must pass the path to a pcap file to return a summary of packet information and traffic data""" 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 @@ -97,7 +107,7 @@ def get_wireless_interface(params: str) -> str: @tool 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}" response = requests.get(url) if response.status_code == 200: