Add functionality to DbManager + WinModuleHashNode refactor to WinPageHashNode

Author: danielhuici

apotheosis.py

@@ -147,14 +147,14 @@ def _create_empty(self, M=0, ef=0, Mmax=0, Mmax0=0,\
                         heuristic=False, extend_candidates=True, keep_pruned_conns=True,\
                         beer_factor: float=0):
         # check first if algorithm is supported
-        if not issubclass(distance_algorithm, HashAlgorithm):
+        if not isinstance(distance_algorithm, type) or not issubclass(distance_algorithm, HashAlgorithm):
             raise ApotheosisUnsupportedDistanceAlgorithmError
-        # construct both data structures (a HNSW and a radix tree for all nodes -- will contain @WinModuleHashNode)
+        # construct both data structures (a HNSW and a radix tree for all nodes -- will contain @WinPageHashNode)
         self._HNSW = HNSW(M=M, ef=ef, Mmax=Mmax, Mmax0=Mmax0, distance_algorithm=distance_algorithm,\
                                 heuristic=heuristic, extend_candidates=extend_candidates, keep_pruned_conns=keep_pruned_conns,\
                                 beer_factor=beer_factor)
         self._distance_algorithm = distance_algorithm
-        # radix hash tree for all nodes (of @WinModuleHashNode)
+        # radix hash tree for all nodes (of @WinPageHashNode)
         self._radix = RadixHash(distance_algorithm)
 
     @classmethod
@@ -538,7 +538,7 @@ def knn_search(self, query=None, k:int=0, ef=0, hashid=0):
         """
         if hashid != 0:
             # create node and make the search again...
-            query = WinModuleHashNode(hashid, self.get_distance_algorithm()) 
+            query = WinPageHashNode(hashid, self.get_distance_algorithm()) 
 
         self._sanity_checks(query)
 
@@ -628,7 +628,7 @@ def __eq__(self, other):
 # unit test
 import common.utilities as util
 from datalayer.node.hash_node import HashNode
-from datalayer.node.winmodule_hash_node import WinModuleHashNode
+from datalayer.node.winpage_hash_node import WinPageHashNode
 from datalayer.hash_algorithm.tlsh_algorithm import TLSHHashAlgorithm
 from datalayer.hash_algorithm.ssdeep_algorithm import SSDEEPHashAlgorithm
 from random import random
@@ -672,11 +672,11 @@ def search_knns(apo, query_node):
     hash6 = "T1DF8174A9C2A506FC122292D644816333FEF1B845C419121A0F91CF5359B5B21FA3A305" #fake
     hash7 = "T10381E956C26225F2DAD9D097B381202C62AC793B37082B8A1EACDAC00B37D557E0E714" #fake
 
-    node1 = WinModuleHashNode(hash1, TLSHHashAlgorithm)
-    node2 = WinModuleHashNode(hash2, TLSHHashAlgorithm)
-    node3 = WinModuleHashNode(hash3, TLSHHashAlgorithm)
-    node4 = WinModuleHashNode(hash4, TLSHHashAlgorithm)
-    node5 = WinModuleHashNode(hash5, TLSHHashAlgorithm)
+    node1 = WinPageHashNode(hash1, TLSHHashAlgorithm)
+    node2 = WinPageHashNode(hash2, TLSHHashAlgorithm)
+    node3 = WinPageHashNode(hash3, TLSHHashAlgorithm)
+    node4 = WinPageHashNode(hash4, TLSHHashAlgorithm)
+    node5 = WinPageHashNode(hash5, TLSHHashAlgorithm)
     nodes = [node1, node2, node3]
 
     print("Testing insert ...")

common/utilities.py

@@ -22,14 +22,14 @@ def create_model(npages, nsearch_pages,\
 import time
 import datetime
 from common.errors import NodeAlreadyExistsError
-def load_DB_in_model(npages=0, nsearch_pages=None, algorithm=None, current_model=None, printlog=True):
+def load_DB_in_model(npages=0, nsearch_pages=None, algorithm=None, current_model=None, printlog=True, modules_of_interest=None, os_id=None):
     BATCH_PRINT=1e2
 
     db_manager = DBManager()
 
     print(f"[*] Getting modules from DB (with {algorithm.__name__}) ...")
     start = time.time_ns()
-    all_pages = db_manager.get_winmodules(algorithm, npages + nsearch_pages if nsearch_pages else npages)
+    all_pages = db_manager.get_winmodules(algorithm, npages + nsearch_pages if nsearch_pages else npages, modules_of_interest, os_id)
     end = time.time_ns() # in nanoseconds
     db_time = (end - start)/1e6 # ms
     print(f"[*] {len(all_pages)} pages recovered from DB in {db_time} ms.")

datalayer/db_manager.py

@@ -5,7 +5,7 @@
 import mysql.connector
 from mysql.connector import errorcode
 
-from datalayer.node.winmodule_hash_node import WinModuleHashNode
+from datalayer.node.winpage_hash_node import WinPageHashNode
 from datalayer.hash_algorithm.hash_algorithm import HashAlgorithm
 from datalayer.hash_algorithm.tlsh_algorithm import TLSHHashAlgorithm
 from datalayer.hash_algorithm.ssdeep_algorithm import SSDEEPHashAlgorithm
@@ -14,6 +14,18 @@
 
 from common.errors import HashValueNotInDBError, PageIdValueNotInDBError
 
+SQL_GET_ALL_OS = """
+    SELECT *
+    FROM os
+    """
+SQL_GET_MODULES_BY_OS = """
+    SELECT m.id AS module_id, m.file_version, m.original_filename, 
+        m.internal_filename, m.product_filename, m.company_name, 
+        m.legal_copyright, m.classification, m.size, m.base_address
+    FROM modules m
+    WHERE m.os_id = %s
+    """
+
 SQL_GET_ALL_PAGES = """
     SELECT p.{}, m.id AS module_id, m.file_version, m.original_filename, 
         m.internal_filename, m.product_filename, m.company_name, 
@@ -78,9 +90,11 @@ def _clean_dict_keys(self, _dict: dict, keys: list):
         for key in keys:
             _dict.pop(key, None)
 
-    def _row_to_module(self, row):
+    def _row_to_module(self, row, os=None):
+        if not os:
+            os = OS(row["id"], row["name"], row["version"])
         return  Module(
-            os=OS(row["id"], row["name"], row["version"]),
+            os=os,
             id=row["module_id"],
             file_version=row["file_version"],
             original_filename=row["original_filename"],
@@ -105,7 +119,7 @@ def get_winmodule_data_by_pageid(self, page_id=0, algorithm=HashAlgorithm):
             raise PageIdValueNotInDBError
 
         module = self._row_to_module(row)
-        return WinModuleHashNode(id=row[hash_column], hash_algorithm=algorithm, module=module)
+        return WinPageHashNode(id=row[hash_column], hash_algorithm=algorithm, module=module)
 
 
     def get_winmodule_data_by_hash(self, algorithm: str = "", hash_value: str = ""):
@@ -119,20 +133,63 @@ def get_winmodule_data_by_hash(self, algorithm: str = "", hash_value: str = ""):
             raise HashValueNotInDBError
 
         module = self._row_to_module(row)
-        return WinModuleHashNode(id=hash_value, hash_algorithm=algorithm, module=module)
-
+        return WinPageHashNode(id=hash_value, hash_algorithm=algorithm, module=module)
 
-    def get_winmodules(self, algorithm: HashAlgorithm = TLSHHashAlgorithm, limit: int = None, modules_of_interest: set = None) -> set:
+    
+    def get_organized_modules(self, algorithm: HashAlgorithm = TLSHHashAlgorithm) -> dict:
+        result = {}
+
+        self.cursor.execute(SQL_GET_ALL_OS)
+        db_operating_systems = self.cursor.fetchall()
+        for db_os in db_operating_systems:
+            os_name = db_os['name']
+            result[os_name] = {}
+
+            self.cursor.execute(SQL_GET_MODULES_BY_OS, (db_os['id'],))
+            modules = self.cursor.fetchall()
+            for module in modules:
+                module_name = module['internal_filename']
+                result[os_name][module_name] = set()
+
+                pages = self.get_winmodules(algorithm, modules_of_interest={module_name}, os_id=db_os['id'])
+                for page in pages:
+                    result[os_name][module_name].add(page)
+
+        return result
+    
+    def get_winmodules(self, algorithm: HashAlgorithm = TLSHHashAlgorithm, limit: int = None, modules_of_interest: set = None, os_id: int = None) -> set:
         try:
             winmodules = set()
             operating_systems = set()
             modules = set()
 
             hash_column = "hashTLSH" if algorithm == TLSHHashAlgorithm else "hashSSDEEP"
-            query = SQL_GET_ALL_PAGES.format(hash_column)
+            query = SQL_GET_ALL_PAGES.format(hash_column)  # Inject hash column
+
+            conditions = []
+            params = []
+
+            if modules_of_interest:
+                placeholders = ', '.join(['%s'] * len(modules_of_interest))
+                conditions.append(f"m.internal_filename IN ({placeholders})")
+                params.extend(modules_of_interest)
+
+            if os_id is not None:
+                conditions.append("o.id = %s")
+                params.append(os_id)
+
+            if algorithm == TLSHHashAlgorithm:
+                conditions.append("p.hashTLSH != '*'")
+                conditions.append("p.hashTLSH != '-'")
+
+            if conditions:
+                query += " WHERE " + " AND ".join(conditions)
+
             if limit:
-                query = query + f" LIMIT {limit}"
-            self.cursor.execute(query)
+                query += " LIMIT %s"
+                params.append(limit)
+
+            self.cursor.execute(query, params)
             results = self.cursor.fetchall()
 
             for row in results:
@@ -148,26 +205,15 @@ def get_winmodules(self, algorithm: HashAlgorithm = TLSHHashAlgorithm, limit: in
                     operating_systems.add(current_os)
 
                 current_module = self._row_to_module(row)
-
-                # Supposedly more memory-efficient, but it slows down retrieval
-                '''
-                if current_module in modules:
-                    current_module = next(module for module in modules if module == current_module)
-                else:
-                    modules.add(current_module)
-                '''
-
-                if modules_of_interest and current_module.internal_filename not in modules_of_interest:
-                    continue
-                
-                winmodules.add(WinModuleHashNode(hash_value, algorithm, current_module))
+                winmodules.add(WinPageHashNode(hash_value, algorithm, current_module))
 
             return winmodules
         except mysql.connector.Error as err:
             logger.error(f"Database query error: {err}")
             raise
         finally:
-            self.cursor.close()
+            pass
+            #self.cursor.close()
 
     def close(self):
         self.cursor.close()

datalayer/node/node.py

@@ -40,7 +40,7 @@ def set_neighbors_at_layer(self, layer: int, neighbors: set):
     # only in HashNode
     def calculate_similarity(self, other_node):
         raise NotImplementedError
-    # only in WinModuleHashNode
+    # only in WinPageHashNode
     def get_pageids(self):
         raise NotImplementedError
 

datalayer/node/winmodule_hash_node.py datalayer/node/winpage_hash_node.pydatalayer/node/winmodule_hash_node.py renamed to datalayer/node/winpage_hash_node.py

@@ -7,7 +7,7 @@
 from common.constants import *
 from common.errors import NodeUnsupportedAlgorithm
 
-class WinModuleHashNode(HashNode):
+class WinPageHashNode(HashNode):
     def __init__(self, id, hash_algorithm: HashAlgorithm, module: Module=None):
         super().__init__(id, hash_algorithm)
         self._module = module
@@ -78,7 +78,7 @@ def create_node_from_DB(cls, db_manager, hash_id, hash_algorithm):
     @classmethod
     def internal_data_needs_DB(cls) -> bool:
         return True # we have some data necessary to retrieve from the DB
-                    # to load a WinModuleHashNode from an Apotheosis file
+                    # to load a WinPageHashNode from an Apotheosis file
 
     def is_equal(self, other):
         if type(self) != type(other):

rest.py

@@ -143,7 +143,7 @@ def _extend_results_winmodule_data(hash_algorithm: str, results: dict) -> dict:
     """Extends the results dict with Winmodule information (from the database).
 
     Arguments:
-    results -- dict of WinModuleHashNode
+    results -- dict of WinPageHashNode
     """
 
     new_results = {}