From 4d86bb01f4ca0c5c5e46dae766ffb9e5923cc2f8 Mon Sep 17 00:00:00 2001 From: David Westgate Date: Fri, 14 Jun 2024 00:57:48 -0700 Subject: [PATCH] work on final --- .gitignore | 1 + final/.env.local | 7 +++ final/README.md | 9 +++ final/app.py | 14 ++--- final/tools.py | 150 +++++++++++++++++++++++++++++++++++++---------- 5 files changed, 141 insertions(+), 40 deletions(-) create mode 100644 final/.env.local create mode 100644 final/README.md diff --git a/.gitignore b/.gitignore index cd77629..71a9470 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ rag_data *wordlist* *dumps* .aider* +*pcap \ No newline at end of file diff --git a/final/.env.local b/final/.env.local new file mode 100644 index 0000000..c8ee55f --- /dev/null +++ b/final/.env.local @@ -0,0 +1,7 @@ +GOOGLE_API_KEY="" +HUGGINGFACEHUB_API_TOKEN="" +SERPAPI_API_KEY="" +VIRUSTOTAL_API_KEY="" +OPENAI_API_KEY="" +FIREWALL_API_KEY="" +LANGCHAIN_API_KEY="" diff --git a/final/README.md b/final/README.md new file mode 100644 index 0000000..5af5bee --- /dev/null +++ b/final/README.md @@ -0,0 +1,9 @@ +###### David Westgate 14 June 2024 +## Final Project for gensec + + +## Prerequisites + + +### Setup + Run + diff --git a/final/app.py b/final/app.py index 537d327..2673cbb 100644 --- a/final/app.py +++ b/final/app.py @@ -3,11 +3,11 @@ from langchain_openai import ChatOpenAI from langchain.agents import AgentExecutor, create_react_agent from langchain.tools import tool from tools import ( + dalle, + tcp_dump, + tshark, get_wireless_interface, - change_adapter_mode, - wifi_encryption_cracking, - deauth_and_capture, - reconnaissance, + ip_loc ) from langchain import hub from langchain_community.tools import ShellTool @@ -28,11 +28,7 @@ client = Client() shell_tool = ShellTool() llm = ChatOpenAI(model_name="gpt-4o", temperature=0) tools = [] -tools.extend( - [ - - ] -) +tools.extend([tcp_dump,ip_loc, dalle,get_wireless_interface]) base_prompt = hub.pull("langchain-ai/react-agent-template") prompt = base_prompt.partial( instructions=""" diff --git a/final/tools.py b/final/tools.py index 1c9abc5..ac491b1 100644 --- a/final/tools.py +++ b/final/tools.py @@ -1,12 +1,19 @@ +from datetime import date, datetime import json import threading import random +from pyshark import FileCapture, packet as Packet from langchain_core.pydantic_v1 import BaseModel, Field from langchain.tools import tool from langchain_community.tools import ShellTool from langchain.chains import LLMChain from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper -from scapy import *; +from scapy import * +import requests +from scapy.all import rdpcap, wrpcap +from scapy.layers.inet import IP, TCP, UDP + + """ """ @@ -20,53 +27,134 @@ shell_tool = ShellTool() class TCPDump(BaseModel): params: str = Field( - description="""A string of all parameters to the `tcpdump` command, including arguments and flags""" + description=""" + A JSON string including `interface` the network interface name and `flags`, any other additional flags or arguments for tcpdump. + Do not include any indication of duration or output file, as this is handled by the tool + + """ ) @tool( - "Perform packet analysis using tcpdump", + "Perform packet capture on a particular network interface using tcpdump. This will return a pcap file which can be analyzed by other tools", args_schema=TCPDump, - return_direct=True, + return_direct=False, ) -def tcp_dump(params: str) -> str: - """Must pass all parameters to `tcpdump` including arguments and flags to perform packet analysis""" - res = shell_tool.run({"commands": [f"tcpdump {params}"]}) - return res +def tcp_dump(json_params: str) -> str: + """Must pass parameters to `tcpdump` including arguments and flags to perform packet monitoring and analysis""" + json_params = json_params.replace("\\", "").replace("`", "").replace("'", "") + print("params ", json_params) + json_obj = json.loads(json_params) + interface = json_obj.get("interface", "") + flags = json_obj.get("flags", "") + # print('date ',date()) + file_name = f'captures/{interface}-{datetime.now().strftime("%Y-%m-%d-%H-%M")}.pcap' + shell_tool.run( + { + "commands": [ + f"sudo timeout 10 tcpdump -i {interface} -w {file_name} {flags} " + ] + } + ) + filtered = filter_unique_packets(file_name) + return filtered class TShark(BaseModel): params: str = Field( - description="""A string of all parameters to the `tshark` command, including arguments and flags""" + description="""A JSON string including `file` the name of the pcap file to analysize, and `flags` and flags any other additional flags or arguments for tshark""" ) + + @tool( - "Perform packet analysis using tshark", + "Perform packet analysis of a pcap file tshark", args_schema=TShark, - return_direct=True, + return_direct=False, ) -def tshark(params: str) -> str: +def tshark(json_params: str) -> str: """Must pass all parameters to `tshark` including arguments and flags to perform packet analysis""" - res = shell_tool.run({"commands": [f"tcpdump {params}"]}) + json_params = json_params.replace("\\", "").replace("`", "").replace("'", "") + print("params ", json_params) + json_obj = json.loads(json_params) + file = json_obj.get("file", "") + flags = json_obj.get("flags", "") + + +class Image(BaseModel): + params: str = Field( + description="""A Graphviz decription of network characteristics""" + ) + + +@tool( + "Create an image of a network", + args_schema=TShark, + return_direct=False, +) +def dalle(params: str) -> str: + """Must provive a Graphviz decription of network characteristics to pass to an image generation tool""" + prompt = f""" + Use the below Graphviz description to create an image of network topology + ``` + {params} + ``` + + """ + image_url = DallEAPIWrapper().run(prompt) + return image_url + + +# From hw6 +class Iwconfig(BaseModel): + params: str = Field( + description="should be command line parameters to 'iwconfig'. If none are needed, this should be left as an empty string" + ) + + +@tool("Get interface information", args_schema=Iwconfig, return_direct=False) +def get_wireless_interface(params: str) -> str: + """Return wireless interface information via iwconfig""" + params = params.replace("`", "").replace("\n", "") # fix buggy input from LLM + print("params ", params) + res = shell_tool.run({"commands": [f"iwconfig {params}"]}) return res -# TODO: Scapy tool +@tool +def ip_loc(address): + """Get information from an ip address, including geolocation""" + url = f"http://ipwho.is/{address}" + response = requests.get(url) + if response.status_code == 200: + return response.json() -# @tool( -# "Perform wifi encryption cracking with aircrack-ng", -# args_schema=CrackPassword, -# return_direct=True, -# ) -# def wifi_encryption_cracking(json_params: str) -> str: -# """Must pass bssid and capfile parameters (as a string containing a json object) to aircrack-ng to perform wifi encryption cracking""" -# json_params = json_params.replace("\\", "") -# json_obj = json.loads(json_params) -# bssid = json_obj["bssid"] -# cap = json_obj["capfile"] -# res = shell_tool.run( -# { -# "commands": [f"aircrack-ng --bssid {bssid} -w wordlist/rockyou.txt {cap}"] -# } # TODO: Abstrace out wordlist - allow custom/multiple wordlist files as long as they sit in the directory -# ) -# return res +# Filter unique packets from a pcap file with pyshark +def filter_unique_packets(pcap_file: str): + packets = rdpcap(pcap_file) + unique_packets = set() + filtered_packets = [] + + for packet in packets: + try: + if IP in packet and (TCP in packet or UDP in packet): + ip_layer = packet[IP] + src_ip = ip_layer.src + dst_ip = ip_layer.dst + + if TCP in packet: + protocol = "TCP" + src_port = packet[TCP].sport + dst_port = packet[TCP].dport + elif UDP in packet: + protocol = "UDP" + src_port = packet[UDP].sport + dst_port = packet[UDP].dport + + packet_signature = (src_ip, dst_ip, protocol, src_port, dst_port) + + if packet_signature not in unique_packets: + unique_packets.add(packet_signature) + filtered_packets.append(packet) + except AttributeError: + continue