Merge pull request #263 from nyu-mlab/packets

Upgrade API fields and Debug Packet Timestamps

Author: crazyideas21

pyproject.toml

@@ -9,7 +9,7 @@ description = "IoT Inspector Client for analyzing IoT device firmware"
 readme = "README.md"
 requires-python = "==3.13.*"
 dependencies = [
-    "libinspector==1.0.10",
+    "libinspector==1.0.12",
     "streamlit>=1.50.0",
 ]
 authors = [

src/libinspector/common.py

@@ -4,7 +4,7 @@
 import json
 import functools
 import pandas as pd
-from typing import Any
+import typing
 import streamlit as st
 import logging
 import re
@@ -157,15 +157,14 @@ def plot_traffic_volume(df: pd.DataFrame, now: int, chart_title: str):
         st.bar_chart(df.set_index('seconds_ago')['Bits'], width='content')
 
 
-def get_device_metadata(mac_address: str):
+def get_device_metadata(mac_address: str) -> dict:
     """
     Retrieve the DHCP hostname and OUI vendor for a device from the database.
 
     Args:
         mac_address (str): The MAC address of the device.
-
     Returns:
-        tuple: (dhcp_hostname, oui_vendor) as strings. Returns empty strings if not found.
+        dict: A dictionary containing the device's metadata, or an empty dictionary if not found.
     """
     db_conn, rwlock = libinspector.global_state.db_conn_and_lock
     sql = """
@@ -175,13 +174,11 @@ def get_device_metadata(mac_address: str):
     with rwlock:
         row = db_conn.execute(sql, (mac_address,)).fetchone()
         if row:
-            metadata = json.loads(row['metadata_json'])
-            dhcp_hostname = metadata.get('dhcp_hostname', "")
-            oui_vendor = metadata.get('oui_vendor', "")
+            meta_data = json.loads(row['metadata_json'])
+            logger.info(f"Parsed Metadata: {json.dumps(meta_data, indent=4)}")  # Check the parsed dictionary
+            return meta_data
         else:
-            dhcp_hostname = ""
-            oui_vendor = ""
-    return dhcp_hostname, oui_vendor
+            return dict()
 
 
 def get_remote_hostnames(mac_address: str):
@@ -247,7 +244,7 @@ def initialize_config_dict():
     config_dict['app_start_time'] = time.time()
 
 
-def config_get(key, default=None) -> Any:
+def config_get(key, default=None) -> typing.Any:
     """
     Get a configuration value.
 
@@ -289,7 +286,7 @@ def config_get_prefix(key_prefix: str):
         }
 
 
-def config_set(key: str, value: Any):
+def config_set(key: str, value: typing.Any):
     """
     Set a configuration value.
 

src/libinspector/device_detail_page.py

@@ -41,23 +41,26 @@ def worker_thread():
 
         # Getting inputs and calling API
         for device_dict in device_list:
-            dhcp_hostname, oui_vendor = common.get_device_metadata(device_dict['mac_address'])
+            meta_data = common.get_device_metadata(device_dict['mac_address'])
             remote_hostnames = common.get_remote_hostnames(device_dict['mac_address'])
             try:
-                api_output = call_predict_api(dhcp_hostname, oui_vendor, remote_hostnames, device_dict['mac_address'])
+                # Note I am passing the metadata as a string because functions with cache cannot take dicts
+                # as a dict is mutable, and the cache would not work as expected.
+                api_output = call_predict_api(json.dumps(meta_data), remote_hostnames, device_dict['mac_address'])
                 common.config_set(f'device_details@{device_dict["mac_address"]}', api_output)
                 if "Vendor" in api_output:
                     # Update name based on API
                     custom_name_key = f"device_custom_name_{device_dict['mac_address']}"
                     custom_name = api_output["Vendor"]
                     if api_output["Vendor"] != "":
                         common.config_set(custom_name_key, custom_name)
-            except RuntimeError:
+            except Exception as e:
+                logger.info("[Device ID API] Exception when calling API: %s", str(e))
                 continue
 
 
 @functools.cache
-def call_predict_api(dhcp_hostname: str, oui_vendor: str, remote_hostnames: str,
+def call_predict_api(meta_data_string: str, remote_hostnames: str,
                      mac_address: str, url="https://dev-id-1.tailcedbd.ts.net/predict") -> dict:
     """
     Call the predicting API with the given fields.
@@ -69,8 +72,7 @@ def call_predict_api(dhcp_hostname: str, oui_vendor: str, remote_hostnames: str,
     2. dhcp_hostname: this is extracted from the 'devices' table, check meta-data and look for 'dhcp_hostname' key.
     3. remote_hostnames: IoT Inspector collects this information the DHCP hostname via either DNS or SNI
     Args:
-        dhcp_hostname (str): The DHCP hostname of the device we want to use AI to get more info about
-        oui_vendor (str): The OUI vendor of the device we want to use AI to get more info about
+        meta_data_string (str): Device Metadata, User Agent info, OUI info, DHCP hostname, etc. in string format
         remote_hostnames (str): The remote hostnames the device has contacted
         mac_address (str): The MAC address of the device we want to use AI to get more info about
         url (str): The API endpoint.
@@ -79,6 +81,7 @@ def call_predict_api(dhcp_hostname: str, oui_vendor: str, remote_hostnames: str,
     """
     api_key = os.environ.get("API_KEY", "momo")
     device_tracked_key = f'tracked@{mac_address}'
+    meta_data = json.loads(meta_data_string)
 
     headers = {
         "Content-Type": "application/json",
@@ -88,11 +91,12 @@ def call_predict_api(dhcp_hostname: str, oui_vendor: str, remote_hostnames: str,
         "prolific_id": common.config_get("prolific_id", ""),
         "mac_address": mac_address,
         "fields": {
-            "oui_friendly": oui_vendor,
-            "dhcp_hostname": dhcp_hostname,
+            "oui_friendly": meta_data.get("oui_vendor", ""),
+            "dhcp_hostname": meta_data.get("dhcp_hostname", ""),
             "remote_hostnames": remote_hostnames,
-            "user_agent_info": "",
-            "netdisco_info": "",
+            "user_agent_info": meta_data.get("user_agent_info", ""),
+            "mdns_info": meta_data.get("mdns_json", ""),
+            "ssdp_info": meta_data.get("ssdp_json", ""),
             "user_labels": "",
             "talks_to_ads": common.config_get(device_tracked_key, False)
         }

src/libinspector/device_list_page.py

@@ -11,6 +11,7 @@
 import common
 import libinspector.core
 import threading
+import libinspector.global_state
 
 
 def get_page(title, material_icon, show_page_func):
@@ -81,6 +82,14 @@ def start_inspector_once():
             daemon=True,
         )
         api_thread.start()
+        label_thread = threading.Thread(
+            name="Device Label Thread",
+            target=device_detail_page.label_thread,
+            daemon=True
+        )
+        label_thread.start()
+        with libinspector.global_state.global_state_lock:
+            libinspector.global_state.custom_packet_callback_func = device_detail_page.save_labeled_activity_packets
 
 
 device_list_page_obj = get_page(

src/libinspector/page_manager.py

@@ -32,20 +32,20 @@
 # and prevent the Flask server from starting with a bad configuration.
 try:
     # Attempt a simple database operation (list collection names) to force connection/auth check
-    print("Attempting connection and write test to MongoDB...")
+    app.logger.info("Attempting connection and write test to MongoDB...")
     # 1. Attempt to insert a test document
     db.test_connection.insert_one({"status": "startup_check", "timestamp": int(time.time())})
     # 2. Attempt to delete the test document
     db.test_connection.delete_one({"status": "startup_check"})
     # 3. Final check to ensure we can list collections
     db.list_collection_names()
-    print("Successfully connected and confirmed write access to MongoDB.")
+    app.logger.info("Successfully connected and confirmed write access to MongoDB.")
 except Exception as e:
-    print("=" * 50)
-    print("FATAL ERROR: Could not connect to MongoDB or authentication failed.")
-    print(f"Connection URI: {mongo_uri}")
-    print(f"Exception: {e}")
-    print("=" * 50)
+    app.logger.info("=" * 50)
+    app.logger.info("FATAL ERROR: Could not connect to MongoDB or authentication failed.")
+    app.logger.info(f"Connection URI: {mongo_uri}")
+    app.logger.info(f"Exception: {e}")
+    app.logger.info("=" * 50)
     sys.exit(1)
     # If using gunicorn/uwsgi, this exit won't stop the workers, but will prevent the app from running correctly.
     # If running with 'python app.py', this will stop the application.
@@ -101,22 +101,40 @@ def label_packets():
         400: Error if required fields are missing or packet decoding fails.
         500: Error if database insertion fails.
     """
+    raw_packets = []
+    capture_times = []
+    capture_lengths = []
+
     data = request.get_json()
-    print("Received POST data:", json.dumps(data, indent=4))
+    app.logger.info("Received POST data:", json.dumps(data, indent=4))
     required_keys = ["packets", "prolific_id", "mac_address", "device_name", "activity_label", "start_time", "end_time"]
     if not data or not all(key in data for key in required_keys):
+        app.logger.warning("Missing required fields in POST data")
         return jsonify({"error": "Missing required fields"}), 400
+
+    if not is_prolific_id_valid(data["prolific_id"]):
+        app.logger.warning("Invalid Prolific ID received")
+        return jsonify({"error": "Prolific ID is invalid"}), 500
+
     try:
-        raw_packets = [base64.b64decode(pkt) for pkt in data["packets"]]
+        for pkt_metadata in data["packets"]:
+            # Validate essential keys are present
+            if not isinstance(pkt_metadata, dict) or 'time' not in pkt_metadata or 'raw_data' not in pkt_metadata:
+                app.logger.error("Packet object missing 'time' or 'raw_data' key.")
+                return jsonify({"error": "Packet metadata structure is invalid"}), 400
+
+            raw_data_bytes = base64.b64decode(pkt_metadata["raw_data"])
+            raw_packets.append(raw_data_bytes)
+            capture_times.append(float(pkt_metadata["time"]))
+            capture_lengths.append(len(raw_data_bytes))
     except Exception as e:
-        print(f"Packet decoding occurred for collection '{data['prolific_id']}': {e}")
+        app.logger.warning(f"Packet decoding occurred for collection '{data['prolific_id']}': {e}")
         return jsonify({"error": "Packet decoding failed"}), 400
 
-    if not is_prolific_id_valid(data["prolific_id"]):
-        return jsonify({"error": "Prolific ID is invalid"}), 500
     folder_path: str = os.path.join(str(data["prolific_id"]), str(data["device_name"]), str(data["activity_label"]))
     fullpath = os.path.normpath(os.path.join(packet_root_dir, folder_path))
     if not fullpath.startswith(packet_root_dir):
+        app.logger.warning("Invalid characters detected in path components")
         return jsonify({"error": "Seems like invalid characters used in prolific ID, device name or activity label"}), 500
 
     prolific_user_packets_collected = db[data["prolific_id"]]
@@ -127,22 +145,24 @@ def label_packets():
         "activity_label": data["activity_label"],
         "start_time": int(data["start_time"]),
         "end_time": int(data["end_time"]),
-        "raw_packets": raw_packets
+        "raw_packets": raw_packets,
+        "capture_times": capture_times,
+        "capture_lengths": capture_lengths
     }
     try:
         prolific_user_packets_collected.insert_one(doc)
     except Exception as e:
-        print(f"MongoDB Insert FAILED for collection '{data['prolific_id']}': {e}")
+        app.logger.warning(f"MongoDB Insert FAILED for collection '{data['prolific_id']}': {e}")
         return jsonify({"error": "Database insert failed"}), 500
 
     try:
         os.makedirs(fullpath, exist_ok=True)
         pcap_file_name: str = make_pcap_filename(int(data["start_time"]), int(data["end_time"]))
         pcap_name: str = os.path.join(fullpath, pcap_file_name)
-        save_packets_to_pcap(raw_packets, pcap_name)
+        save_packets_to_pcap(raw_packets, capture_times, pcap_name)
     except Exception as e:
         # If file saving fails, return a 500 but note that the DB save succeeded
-        print(f"PCAP File Save FAILED for ID: {e}")
+        app.logger.warning(f"PCAP File Save FAILED for ID: {e}")
         return jsonify({
             "status": "partial_success",
             "inserted": 1,
@@ -170,22 +190,27 @@ def make_pcap_filename(start_time: int, end_time: int) -> str:
     return filename
 
 
-def save_packets_to_pcap(raw_packets: list, filename="output.pcap"):
+def save_packets_to_pcap(raw_packets: list, capture_times: list, filename="output.pcap"):
     """
     Saves a list of raw packet bytes to a pcap file.
 
     Args:
         raw_packets (list): List of bytes objects representing raw packets.
+        capture_times (list): The epoch time of each packet when it was captured by IoT Inspector.
         filename (str): Output pcap file name.
     """
     scapy_packets = []
-    for pkt_bytes in raw_packets:
+    for i, pkt_bytes in enumerate(raw_packets):
         try:
             pkt = Ether(pkt_bytes)
             if pkt.__class__.__name__ == "Raw":
                 pkt = IP(pkt_bytes)
         except Exception:
             pkt = IP(pkt_bytes)
+
+        # CRITICAL STEP: Assign the supplied original capture time
+        # This is guaranteed to be correct for ALL packets.
+        pkt.time = capture_times[i]
         scapy_packets.append(pkt)
     wrpcap(filename, scapy_packets)