ail-framework/tools/reprocess_objects.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Reprocess AIL Objects by Object Type
================

Send ALL objects by type in queues

"""

import argparse
import os
import sys

sys.path.append(os.environ['AIL_BIN'])
##################################
# Import Project packages
##################################
from lib.ail_core import is_object_type
from lib import ail_queues
from lib.objects import ail_objects

# from modules.ApiKey import ApiKey
# from modules.Categ import Categ
# from modules.CreditCards import CreditCards
# 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.CodeReader import CodeReader
from modules.Languages import Languages
from modules.OcrExtractor import OcrExtractor

MODULES = {
    'CEDetector': CEDetector,
    'CodeReader': CodeReader,
    'Languages': Languages,
    'OcrExtractor': OcrExtractor
}

def reprocess_message_objects(object_type, module_name=None):
    if module_name:
        module = MODULES[module_name]()
        for obj in ail_objects.obj_iterator(object_type, 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={}):
            queue.send_message(obj.get_global_id(), message='reprocess')
        queue.end()


if __name__ == "__main__":

    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')

    args = parser.parse_args()
    if not args.type:
        parser.print_help()
        sys.exit(0)

    obj_type = args.type
    if not is_object_type(obj_type):
        raise Exception(f'Invalid Object Type: {obj_type}')
    if obj_type not in ['image', 'item', 'message', 'screenshot', 'title']:
        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)