Full migration from orm to sql

Author: danielhuici

common/utilities.py

@@ -29,32 +29,32 @@ def load_DB_in_model(npages=0, nsearch_pages=None, algorithm=None, current_model
 
     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)
     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.")
 
     page_list = []
     insert_times = []
-    for i in range(0, len(all_pages[:npages])):
+    for i, winmodule in enumerate(0, 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()}) ...")
 
         try:
             start = time.time_ns()
-            current_model.insert(all_pages[i]) # can raise exception
+            current_model.insert(winmodule) # can raise exception
             end = time.time_ns() # in nanoseconds
             insert_times.append((end - start)/(1e6)) # convert to ms
-            page_list.append(all_pages[i].get_id())
+            page_list.append(winmodule.get_id())
         except NodeAlreadyExistsError: # it should never occur...
             # get module already in DB, and print it to compare with the other one
-            exact, node = current_model.search_exact_match_only(all_pages[i].get_id())
+            exact, node = current_model.search_exact_match_only(winmodule.get_id())
             if not exact: # infeasible path. If you see this, something weird happened
                 raise Exception # db was modified in the backend, don't worry ...
 
             # check they are _really_ the same
             existing_page = node.get_page()
-            new_page = all_pages[i].get_page()
+            new_page = winmodule.get_page()
             equal, equal_test = pages_are_equal(existing_page, new_page)
             if equal:
                 logging.warning(f"Node \"{node.get_id()}\" already exists (different page id, same hashes)!")

datalayer/database/module.py

@@ -1,32 +1,25 @@
-from sqlalchemy import Column, Integer, String, BigInteger, ForeignKey
-from sqlalchemy.orm import relationship, declarative_base
-from .base import Base
 
-class Module(Base):
-    __tablename__ = 'modules'
+class Module:
+    def __init__(self, os, id, file_version, original_filename,
+        internal_filename, product_name,company_name, legal_copyright,
+        classification, size, base_address
+        ):
 
-    # April 05, 2024: Updated for dataset DB
-    id                = Column(BigInteger, primary_key=True)
-    file_version      = Column(String)
-    original_filename = Column(String)
-    internal_filename = Column(String)
-    product_name      = Column(String)
-    company_name      = Column(String)
-    legal_copyright   = Column(String)
-    classification    = Column(String)
-    size              = Column(Integer)
-    base_address      = Column(BigInteger)
-    cpu               = Column(String)
-    os_id             = Column(BigInteger, ForeignKey('os.id'))
+        self.id = id
+        self.file_version = file_version
+        self.original_filename = original_filename
+        self.internal_filename = internal_filename
+        self.product_name = product_name
+        self.company_name = company_name
+        self.legal_copyright = legal_copyright
+        self.classification = classification
+        self.size = size
+        self.base_address = base_address
+        self.os = os
 
-    pages = relationship("Page")
-    os = relationship("OS", back_populates="modules")
-
-    def as_dict(self):
-       return {c.name: getattr(self, c.name) for c in self.__table__.columns} #| self.os.as_dict()
-
-    def __str__(self):
-        return str(self.as_dict())
+    def __eq__(self, module):
+        return self.id == module.id
 
-    def __repr__(self):
-        return str(self.as_dict())
+    def __hash__(self):
+        return hash(self.id)
+        

datalayer/database/operating_system.py

@@ -1,19 +1,20 @@
-from sqlalchemy.orm import relationship, declarative_base
-from sqlalchemy import create_engine, Column, Integer, String, BigInteger
-from .base import Base
 
-class OS(Base):
-    __tablename__ = 'os'
-
-    # April 05, 2024: Updated for dataset DB
-    id      = Column(BigInteger, primary_key=True)
-    name    = Column(String)
-    version = Column(String)
-    cpu     = Column(String)
-    modules = relationship("Module")
+class OS:
+    def __init__(self, id, name, version):
+        self.id = id
+        self.name = name
+        self.version = version
 
     def __str__(self):
         return f"Name: {self.name}, Version: {self.version}"
 
     def as_dict(self):
-       return {str('os_' + c.name): getattr(self, c.name) for c in self.__table__.columns}
+        return {f"os_{key}": value for key, value in self.__dict__.items()}
+
+    def __eq__(self, other):
+        if isinstance(other, OS):
+            return self.id == other.id and self.name == other.name and self.version == other.version
+        return False
+
+    def __hash__(self):
+        return hash((self.id, self.name, self.version))