Minor fixes in datalayer

Author: danielhuici

common/errors.py

@@ -1,5 +1,9 @@
 
 # User-defined exceptions
+class ApotheosisUnsupportedNodeObjectError(Exception):
+    pass
+class ApotheosisUnsupportedDistanceAlgorithmError(Exception):
+    pass
 class ApotheosisUnmatchDistanceAlgorithmError(Exception):
     pass
 class ApotheosisIsEmptyError(Exception):

common/utilities.py

@@ -36,7 +36,7 @@ def load_DB_in_model(npages=0, nsearch_pages=None, algorithm=None, current_model
 
     page_list = []
     insert_times = []
-    for i, winmodule in enumerate(0, all_pages):
+    for i, winmodule in enumerate(all_pages):
         if i % BATCH_PRINT  == 0 and i != 0 and printlog:
             print(f"{int(i/BATCH_PRINT)}*{BATCH_PRINT} pages already inserted ({datetime.datetime.now()}) ...")
 

datalayer/database/module.py

@@ -23,3 +23,17 @@ def __eq__(self, module):
     def __hash__(self):
         return hash(self.id)
 
+    def as_dict(self):
+        return {
+            "id": self.id,
+            "file_version": self.file_version,
+            "original_filename": self.original_filename,
+            "internal_filename": self.internal_filename,
+            "product_name": self.product_name,
+            "company_name": self.company_name,
+            "legal_copyright": self.legal_copyright,
+            "classification": self.classification,
+            "size": self.size,
+            "base_address": self.base_address,
+            "os": self.os.as_dict() if self.os else None
+        }

datalayer/db_manager.py

@@ -131,7 +131,7 @@ def get_winmodules(self, algorithm: HashAlgorithm = TLSHHashAlgorithm, limit: in
             hash_column = "hashTLSH" if algorithm == TLSHHashAlgorithm else "hashSSDEEP"
             query = SQL_GET_ALL_PAGES.format(hash_column)
             if limit:
-                query = SQL_GET_ALL_PAGES + f" LIMIT {limit}"
+                query = query + f" LIMIT {limit}"
             self.cursor.execute(query)
             results = self.cursor.fetchall()
 

datalayer/node/hash_node.py

@@ -40,3 +40,10 @@ def to_JSON(self):
 
     def get_distance_algorithm(self):
         return self._hash_algorithm
+    
+    def as_dict(self):
+        return {
+            "id": self._id,
+            "hash_algorithm": self._hash_algorithm.__name__
+        }
+

datalayer/node/winmodule_hash_node.py

@@ -29,6 +29,30 @@ def get_draw_features(self):
                 "module_version": {self._id: self._module.file_version},
                 "os_version": {self._id: self._module.os.version}
                 }
+    
+
+    def as_dict(self):
+        node_dict = super().as_dict()
+        if self._module:
+            node_dict.update({
+                "module_id": self._module.id,
+                "file_version": self._module.file_version,
+                "original_filename": self._module.original_filename,
+                "internal_filename": self._module.internal_filename,
+                "product_name": self._module.product_name,
+                "company_name": self._module.company_name,
+                "legal_copyright": self._module.legal_copyright,
+                "classification": self._module.classification,
+                "size": self._module.size,
+                "base_address": self._module.base_address,
+            })
+
+            if self._module.os:
+                node_dict.update(self._module.os.as_dict())
+
+        return node_dict
+
+
 
     def internal_serialize(self):
         return self.get_internal_page_id().to_bytes(I_SIZE, byteorder=BYTE_ORDER)