Merge pull request #262 from nyu-mlab/fix-gui

Fix Device API Thread and Labeling Packet Bug

Author: crazyideas21

README.md

@@ -1,28 +1,31 @@
 # IoT Inspector 3
 
-Simply run `./start.bash`. It will take care of all the dependencies.
+Simply run `./start.bash` for Linux/Mac and `start.bat` for Windows. It will take care of all the dependencies.
 
-If the underlying Inspector core library is updated, please run the following first:
+If the underlying dependencies is updated, please run the following first:
 
 ```bash
 uv cache clean
-uv lock --upgrade-package libinspector
+uv lock
 uv sync
 ```
 
+# User guide
+Please review the [User Guide](https://github.com/nyu-mlab/iot-inspector-client/wiki) for instructions how to run IoT Inspector. 
 
 # Developer Guide
 
 If you are developing IoT Inspector, please read this section.
 
 ## Database Schema
 
-When presenting network stats, IoT Inspector reads from an internal SQLite database.
+When presenting network stats, IoT Inspector reads from an internal SQLite database. 
+To see how the packet collector and database is implemented, look at the [IoT Inspector Core package](https://github.com/nyu-mlab/inspector-core-library).
 
 You should always read from the database using the following approach:
 
 ```python
-import libinspector
+import libinspector.global_state
 db_conn, rwlock = libinspector.global_state.db_conn_and_lock
 with rwlock:
     db_conn.execute("SELECT * FROM devices")
@@ -38,15 +41,15 @@ CREATE TABLE devices (
     is_gateway INTEGER DEFAULT 0,
     updated_ts INTEGER DEFAULT 0,
     metadata_json TEXT DEFAULT '{}'
-)
+);
 
 CREATE TABLE hostnames (
     ip_address TEXT PRIMARY KEY,
     hostname TEXT NOT NULL,
     updated_ts INTEGER DEFAULT 0,
     data_source TEXT NOT NULL,
     metadata_json TEXT DEFAULT '{}'
-)
+);
 
 CREATE TABLE network_flows (
     timestamp INTEGER,
@@ -69,5 +72,5 @@ CREATE TABLE network_flows (
             src_port, dest_port,
             protocol
         )
-)
+);
 ```

pyproject.toml

@@ -7,7 +7,7 @@ name = "iot-inspector"
 dynamic = ["version"]
 description = "IoT Inspector Client for analyzing IoT device firmware"
 readme = "README.md"
-requires-python = ">=3.13"
+requires-python = "==3.13.*"
 dependencies = [
     "libinspector==1.0.10",
     "streamlit>=1.50.0",

src/libinspector/common.py

@@ -3,13 +3,13 @@
 import threading
 import json
 import functools
-import libinspector.global_state
-from libinspector.privacy import is_ad_tracked
 import pandas as pd
 from typing import Any
 import streamlit as st
 import logging
 import re
+import libinspector.global_state
+from libinspector.privacy import is_ad_tracked
 
 config_file_name = 'config.json'
 config_lock = threading.Lock()

src/libinspector/device_detail_page.py

@@ -97,7 +97,6 @@ def _send_packets_callback(mac_address: str):
             api_url = f"https://{remote_host}:{remote_port}{api_path}"
 
             # NOTE: We avoid st.write here, save the info to session_state instead
-
             response = requests.post(
                 api_url,
                 json={
@@ -117,15 +116,35 @@ def _send_packets_callback(mac_address: str):
                 st.session_state['api_message'] = f"success| {len(pending_packet_list)} Labeled packets successfully sent to the server."
             else:
                 st.session_state[
-                    'api_message'] = f"error|Failed to send labeled packets. Server status: {response.status_code}."
+                    'api_message'] = (f"error|Failed to send labeled packets. Server status: {response.status_code}. "
+                                      f"{len(pending_packet_list)} Packets were not sent.")
         else:
             st.session_state['api_message'] = "error|No packets were captured for labeling."
-
     except requests.RequestException as e:
         st.session_state['api_message'] = f"error|An error occurred during API transmission: {e}"
+    finally:
+        pending_packet_list.clear()
     # st.rerun() will occur after this, showing the results.
 
 
+def update_device_inspected_status(mac_address: str):
+    """
+    Manually update to inspected status so that all the packets can be collected for the MAC Address.
+    Args:
+        mac_address (str): The MAC address of the device to update.
+    """
+    db_conn, rwlock = libinspector.global_state.db_conn_and_lock
+    with rwlock:
+        sql = """
+            UPDATE devices
+            SET is_inspected = ?
+            WHERE mac_address = ?
+        """
+        db_conn.execute(sql, (1, mac_address))
+    device_inspected_config_key = f'device_is_inspected_{mac_address}'
+    common.config_set(device_inspected_config_key, True)
+
+
 def label_activity_workflow(mac_address: str):
     """
     Manages the interactive, state-driven workflow for labeling network activity in Streamlit.
@@ -145,7 +164,7 @@ def label_activity_workflow(mac_address: str):
     5. Control Buttons: Displays 'Start' and 'Labeling Complete' buttons with dynamic 'disabled' states to enforce the correct sequence.
     6. Countdown: Executes a 5-second blocking countdown in the main thread after 'Start' is clicked and before packet collection begins.
     7. Packet Collection: Starts packet capture by setting a global callback function (save_labeled_activity_packets) for the specified mac_address.
-    8. Status Display: Uses an st.empty() placeholder and the api_message state variable to display feedback immediately beneath the control buttons for superior UX.
+    8. Status Display: Uses a st.empty() placeholder and the api_message state variable to display feedback immediately beneath the control buttons for superior UX.
     9. Final Summary: Upon session completion, displays the recorded activity label, device name, and duration.
 
     Args:
@@ -185,6 +204,7 @@ def label_activity_workflow(mac_address: str):
             disabled=session_active or st.session_state['end_time'] is not None,
             help="Click to start labeling an activity for this device. This will reset any previous labeling state."
     ):
+        update_device_inspected_status(mac_address)
         reset_labeling_state()
         # Keep labeling_in_progress=True until the very end, controlled by config_get/set
         common.config_set('labeling_in_progress', True)

src/libinspector/device_list_page.py

@@ -5,34 +5,25 @@
 import common
 import logging
 import functools
-import threading
 import os
 import requests
 
-
 logger = logging.getLogger("client")
 
+
 def show():
     """
     Creating a page that shows what devices have been discovered so far
     """
     toast_obj = st.toast('Discovering devices...')
     show_list(toast_obj)
 
-@functools.lru_cache(maxsize=1)
-def start_inspector_once():
-    """Initialize the Inspector core only once."""
-    with st.spinner("Starting API Thread..."):
-        threading.Thread(
-            target=worker_thread,
-            daemon=True,
-        )
-
 
 def worker_thread():
     """
     A worker thread to periodically clear the cache of call_predict_api.
     """
+    logger.info("[Device ID API] Starting worker thread to periodically call the API for each device.")
     while True:
         time.sleep(15)
         db_conn, rwlock = libinspector.global_state.db_conn_and_lock
@@ -46,6 +37,7 @@ def worker_thread():
         with rwlock:
             for device_dict in db_conn.execute(sql):
                 device_list.append(dict(device_dict))
+        logger.info("[Device ID API] 15 seconds passed, will start calling API for each device if needed.")
 
         # Getting inputs and calling API
         for device_dict in device_list:
@@ -54,6 +46,12 @@ def worker_thread():
             try:
                 api_output = call_predict_api(dhcp_hostname, oui_vendor, 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:
                 continue
 
@@ -111,8 +109,8 @@ def call_predict_api(dhcp_hostname: str, oui_vendor: str, remote_hostnames: str,
         raise RuntimeError(
             "Fewer than two string fields in data are non-empty; refusing to call API. Wait until IoT Inspector collects more data.")
 
-    # TODO: Used for debugging, risk is minimal since this is client-facing code. Should have some flag for production.
-    logger.info("[Device ID API] Calling API with data: %s", json.dumps(data, indent=4))
+    if common.config_get("debug", default=False):
+        logger.info("[Device ID API] Calling API with data: %s", json.dumps(data, indent=4))
 
     try:
         response = requests.post(url, headers=headers, json=data, timeout=10)

src/libinspector/page_manager.py

@@ -10,6 +10,7 @@
 import functools
 import common
 import libinspector.core
+import threading
 
 
 def get_page(title, material_icon, show_page_func):
@@ -69,7 +70,17 @@ def initialize_page():
 def start_inspector_once():
     """Initialize the Inspector core only once."""
     with st.spinner("Starting Inspector Core Library..."):
+        # Just in case someone closes labeling window without finishing
+        # Same with the general warning
+        common.config_set("suppress_warning", False)
+        common.config_set("labeling_in_progress", False)
         libinspector.core.start_threads()
+        api_thread = threading.Thread(
+            name="Device API Thread",
+            target=device_list_page.worker_thread,
+            daemon=True,
+        )
+        api_thread.start()
 
 
 device_list_page_obj = get_page(