chg: [reprocess tools] reprocess object message by tags

ail-project · Feb 27, 2025 · b668729 · b668729
1 parent a31abb8
commit b668729
Showing 1 changed file with 27 additions and 8 deletions.
diff --git a/tools/reprocess_objects.py b/tools/reprocess_objects.py
@@ -26,33 +26,45 @@
 # from modules.DomClassifier import DomClassifier
 # from modules.Global import Global
 # from modules.Keys import Keys
-# from modules.Onion import Onion
 # from modules.Telegram import Telegram
 
 from modules.CEDetector import CEDetector
+from modules.CveModule import CveModule
 from modules.CodeReader import CodeReader
+from modules.Cryptocurrencies import Cryptocurrencies
+from modules.Decoder import Decoder
 from modules.Languages import Languages
 from modules.OcrExtractor import OcrExtractor
+from modules.Onion import Onion
+from modules.PgpDump import PgpDump
 
 MODULES = {
     'CEDetector': CEDetector,
+    'Cryptocurrencies': Cryptocurrencies,
+    'CveModule': CveModule,
     'CodeReader': CodeReader,
+    'Decoder': Decoder,
     'Languages': Languages,
-    'OcrExtractor': OcrExtractor
+    'OcrExtractor': OcrExtractor,
+    'Onion': Onion,
+    'PgpDump': PgpDump
 }
 
-def reprocess_message_objects(object_type, module_name=None):
+def reprocess_message_objects(object_type, module_name=None, tags=[]):
+    filters = {}
+    if tags:
+        filters['tags'] = tags
     if module_name:
         module = MODULES[module_name]()
-        for obj in ail_objects.obj_iterator(object_type, filters={}):
+        for obj in ail_objects.obj_iterator(object_type, filters=filters):
             if not obj.exists():
                 print(f'ERROR: object does not exist, {obj.id}')
                 continue
             module.obj = obj
             module.compute(None)
     else:
         queue = ail_queues.AILQueue('FeederModuleImporter', -1)
-        for obj in ail_objects.obj_iterator(object_type, filters={}):
+        for obj in ail_objects.obj_iterator(object_type, filters=filters):
             queue.send_message(obj.get_global_id(), message='reprocess')
         queue.end()
 
@@ -62,6 +74,7 @@ def reprocess_message_objects(object_type, module_name=None):
     parser = argparse.ArgumentParser(description='Reprocess AIL Objects')
     parser.add_argument('-t', '--type', type=str, help='AIL Object Type', required=True)
     parser.add_argument('-m', '--module', type=str, help='AIL Module Name')
+    parser.add_argument('--tags', nargs='+', type=str, help='List of tags')
 
     args = parser.parse_args()
     if not args.type:
@@ -75,6 +88,12 @@ def reprocess_message_objects(object_type, module_name=None):
         raise Exception(f'Currently not supported Object Type: {obj_type}')
 
     modulename = args.module
-    if modulename not in MODULES:
-        raise Exception(f'Currently not supported Module: {modulename}')
-    reprocess_message_objects(obj_type, module_name=modulename)
+    if modulename:
+        if modulename not in MODULES:
+            raise Exception(f'Currently not supported Module: {modulename}')
+
+    if args.tags:
+        ltags = args.tags
+    else:
+        ltags = []
+    reprocess_message_objects(obj_type, module_name=modulename, tags=ltags)