Integrate with Device Detection LLM API

Author: AndrewQuijano

common.py

@@ -1,15 +1,112 @@
 import datetime
+import os
 import time
 import threading
 import json
 import functools
-
+import requests
+import libinspector
 
 config_file_name = 'config.json'
 config_lock = threading.Lock()
 config_dict = {}
 
 
+def call_predict_api(mac_address: str, url="https://dev-id-1.tailcedbd.ts.net/predict") -> dict:
+    """
+    Call the predicting API with the given fields.
+    This takes the MAC Address of an inspected device
+    and checks the `devices` table, where iot inspector core collected meta-data
+    based on SSDP discovery.
+
+    Please see Page 11 Table A.1. We explain how to get the data from IoT Inspector:
+    1. oui_friendly: we use the OUI database from IoT Inspector Core
+    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:
+        mac_address (str): The MAC address of the device we want to use AI to get more info about
+        url (str): The API endpoint.
+
+    Returns:
+        dict: The response text from the API.
+    """
+
+    db_conn, rwlock = libinspector.global_state.db_conn_and_lock
+
+    # Get the DHCP Hostname if found
+    sql = """
+        SELECT metadata_json FROM devices
+        WHERE mac_address = ?
+    """
+    with rwlock:
+        row = db_conn.execute(sql, (mac_address,)).fetchone()
+        if row:
+            metadata = json.loads(row['metadata_json'])
+            if 'dhcp_hostname' in metadata:
+                dhcp_hostname = metadata['dhcp_hostname']
+            else:
+                dhcp_hostname = ""
+            if 'oui_vendor' in metadata:
+                oui_vendor = metadata['oui_vendor']
+            else:
+                oui_vendor = ""
+        else:
+            dhcp_hostname = ""
+            oui_vendor = ""
+
+    # Get the remote hostname
+    sql = """
+        SELECT DISTINCT hostname
+        FROM (
+            SELECT src_hostname AS hostname FROM network_flows
+            WHERE src_mac_address = ? AND src_hostname IS NOT NULL
+            UNION
+            SELECT dest_hostname AS hostname FROM network_flows
+            WHERE dest_mac_address = ? AND dest_hostname IS NOT NULL
+        ) AS combined
+    """
+    with rwlock:
+        rows = db_conn.execute(sql, (mac_address, mac_address)).fetchall()
+        hostnames = [row['hostname'] for row in rows if row['hostname']]
+        remote_hostnames = '+'.join(hostnames) if hostnames else ""
+
+    api_key = os.environ.get("API_KEY", "momo")
+    headers = {
+        "Content-Type": "application/json",
+        "x-api-key": api_key
+    }
+    data = {
+        "fields": {
+            "oui_friendly": oui_vendor,
+            "dhcp_hostname": dhcp_hostname,
+            "remote_hostnames": remote_hostnames,
+            "user_agent_info" : "",
+            "netdisco_info" : "",
+            "user_labels" : "",
+            "talks_to_ads" : False
+        }
+    }
+    try:
+        response = requests.post(url, headers=headers, json=data, timeout=10)
+        response.raise_for_status()  # Raises HTTPError for bad responses
+        result = response.json()  # May raise ValueError if not valid JSON
+    except requests.exceptions.RequestException as e:
+        # Handle network errors, timeouts, or HTTP errors
+        print(f"API request failed: {e}")
+        return {}
+    except ValueError as e:
+        # Handle JSON decoding errors
+        print(f"Failed to decode JSON response: {e}")
+        return {}
+
+    # Store the results in the config.json output
+    # The key is MAC Address, and the output is directly written
+    print("API query successful!")
+    config_set(f'device_details@{mac_address}', result)
+    return response.json()
+
+
 def get_human_readable_time(timestamp=None):
     """
     Convert a timestamp to a human-readable time format.
@@ -25,7 +122,6 @@ def get_human_readable_time(timestamp=None):
     return datetime.datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')
 
 
-
 def initialize_config_dict():
     """
     Initialize the configuration dictionary by reading from the config file.
@@ -43,7 +139,6 @@ def initialize_config_dict():
     config_dict['app_start_time'] = time.time()
 
 
-
 def config_get(key, default=None):
     """
     Get a configuration value.
@@ -69,7 +164,6 @@ def config_get(key, default=None):
         raise KeyError(f"Key '{key}' not found in configuration.")
 
 
-
 def config_get_prefix(key_prefix):
     """
     Get all configuration values that start with a given prefix.
@@ -91,7 +185,6 @@ def config_get_prefix(key_prefix):
         }
 
 
-
 def config_set(key, value):
     """
     Set a configuration value.
@@ -131,9 +224,8 @@ def get_device_custom_name(mac_address):
     return device_custom_name
 
 
-
 @functools.cache
-def get_sql_query(sql_file_name):
+def get_sql_query(sql_file_name: str) -> str:
     """
     Get the SQL query from a file.