work on final

.gitignore

@@ -10,3 +10,4 @@ rag_data
 *wordlist*
 *dumps*
 .aider*
+*pcap

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=""

final/README.md

@@ -0,0 +1,9 @@
+###### David Westgate 14 June 2024
+## Final Project for gensec
+
+
+## Prerequisites
+
+
+### Setup + Run
+

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="""

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