server: vulnsearch, versioninfo refactor pruning, match same flow in …

…both components
bodik · Nov 13, 2023 · ee12e57 · ee12e57
1 parent db75c86
commit ee12e57
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 13 deletions.
diff --git a/sner/server/storage/versioninfo.py b/sner/server/storage/versioninfo.py
@@ -101,6 +101,9 @@ def flush(self):
             )
         db.session.commit()
 
+    def prune(self):
+        """prune database of gone items"""
+
         affected_rows = Versioninfo.query.filter(Versioninfo.id.not_in(self.data.keys())).delete(synchronize_session=False)
         current_app.logger.debug('prune versioninfo %d items', affected_rows)
         db.session.commit()
@@ -168,6 +171,7 @@ def rebuild(cls):
         vmap = cls.collect_nmap_mysqlinfo(vmap)
         vmap = cls.collect_nmap_rdpntlminfo(vmap)
         vmap.flush()
+        vmap.prune()
 
     @classmethod
     def collect_nmap_bannerdict(cls, vmap):

diff --git a/sner/server/storage/vulnsearch.py b/sner/server/storage/vulnsearch.py
@@ -106,12 +106,7 @@ def __init__(self, buflen=1000):
 
         self.buf = []
         self.buflen = buflen
-        self.prune_list = []
-
-    def prune_init(self):
-        """initialize prune list"""
-
-        self.prune_list = db.session.execute(select(Vulnsearch.id)).scalars().all()
+        self.prune_list = set()
 
     def index(self, doc_id, doc):
         """index item in buffered way"""
@@ -129,18 +124,15 @@ def flush(self):
                 db.session.execute(
                     pg_insert(Vulnsearch).values(item).on_conflict_do_update(constraint='vulnsearch_pkey', set_=item)
                 )
-                try:
-                    self.prune_list.remove(item["id"])
-                except ValueError:
-                    pass
+                self.prune_list.add(item['id'])
             db.session.commit()
             self.buf = []
 
     def prune(self):
         """prune gone items"""
 
-        current_app.logger.debug('prune vulnsearch %d items', len(self.prune_list))
-        Vulnsearch.query.filter(Vulnsearch.id.in_(self.prune_list)).delete(synchronize_session=False)
+        affected_rows = Vulnsearch.query.filter(Vulnsearch.id.not_in(self.prune_list)).delete(synchronize_session=False)
+        current_app.logger.debug('prune vulnsearch %d items', affected_rows)
         db.session.commit()
         db.session.expire_all()
 
@@ -210,7 +202,6 @@ def rebuild_localdb(self):
         """build local vulnsearch tables"""
 
         vulnsearch_writer = LocaldbWriter(self.rebuild_buflen)
-        vulnsearch_writer.prune_init()
 
         for note, icpe, parsed_cpe in cpe_notes():
             for cve in self.cvefor(icpe):