diff --git a/bin/modules/Global.py b/bin/modules/Global.py index d45f4806..3f473aed 100755 --- a/bin/modules/Global.py +++ b/bin/modules/Global.py @@ -105,6 +105,11 @@ class Global(AbstractModule): filename = self.check_filename(filename, new_file_content) if filename: + new_obj_id = filename.replace(self.ITEMS_FOLDER, '', 1) + new_obj = Item(new_obj_id) + new_obj.sanitize_id() + self.set_obj(new_obj) + # create subdir dirname = os.path.dirname(filename) if not os.path.exists(dirname): diff --git a/bin/modules/Pasties.py b/bin/modules/Pasties.py index 3d420c84..1c74635b 100755 --- a/bin/modules/Pasties.py +++ b/bin/modules/Pasties.py @@ -56,7 +56,7 @@ class Pasties(AbstractModule): with open(domains_pasties) as f: for line in f: url = line.strip() - if url: # TODO validate line + if url: # TODO validate line self.faup.decode(url) url_decoded = self.faup.get() host = url_decoded['host'] @@ -135,7 +135,7 @@ class Pasties(AbstractModule): if path.startswith(url_path): if url_path != path and url_path != path_end: print('send to crawler', url_path, url) - self.send_to_crawler(url, self.obj.id)) + self.send_to_crawler(url, self.obj.id) break diff --git a/bin/modules/abstract_module.py b/bin/modules/abstract_module.py index c22fa568..affd6c6a 100644 --- a/bin/modules/abstract_module.py +++ b/bin/modules/abstract_module.py @@ -76,6 +76,14 @@ class AbstractModule(ABC): def get_obj(self): return self.obj + def set_obj(self, new_obj): + if self.obj: + old_id = self.obj.id + self.obj = new_obj + self.queue.rename_message_obj(self.obj.id, old_id) + else: + self.obj = new_obj + def get_message(self): """ Get message from the Redis Queue (QueueIn) diff --git a/tests/test_modules.py b/tests/test_modules.py index 282376af..19fd04a6 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -10,7 +10,10 @@ from base64 import b64encode from distutils.dir_util import copy_tree sys.path.append(os.environ['AIL_BIN']) - +################################## +# Import Project packages +################################## +from lib.ConfigLoader import ConfigLoader # Modules Classes from modules.ApiKey import ApiKey from modules.Categ import Categ @@ -22,59 +25,58 @@ from modules.Onion import Onion from modules.Telegram import Telegram # project packages -from lib.ConfigLoader import ConfigLoader -import lib.crawlers as crawlers import lib.objects.Items as Items #### COPY SAMPLES #### config_loader = ConfigLoader() -# # TODO:move me in new Item package -ITEMS_FOLDER = os.path.join(os.environ['AIL_HOME'], config_loader.get_config_str("Directories", "pastes")) + '/' -ITEMS_FOLDER = os.path.join(os.path.realpath(ITEMS_FOLDER), '') +ITEMS_FOLDER = Items.ITEMS_FOLDER TESTS_ITEMS_FOLDER = os.path.join(ITEMS_FOLDER, 'tests') sample_dir = os.path.join(os.environ['AIL_HOME'], 'samples') copy_tree(sample_dir, TESTS_ITEMS_FOLDER) #### ---- #### -class Test_Module_ApiKey(unittest.TestCase): +class TestModuleApiKey(unittest.TestCase): def setUp(self): - self.module_obj = ApiKey() - self.module_obj.debug = True + self.module = ApiKey() + self.module.debug = True def test_module(self): item_id = 'tests/2021/01/01/api_keys.gz' + self.module.obj = Items.Item(item_id) google_api_key = 'AIza00000000000000000000000_example-KEY' aws_access_key = 'AKIAIOSFODNN7EXAMPLE' aws_secret_key = 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY' - matches = self.module_obj.compute(f'{item_id} 3', r_result=True) + matches = self.module.compute('3', r_result=True) self.assertCountEqual(matches[0], {google_api_key}) self.assertCountEqual(matches[1], {aws_access_key}) self.assertCountEqual(matches[2], {aws_secret_key}) -class Test_Module_Categ(unittest.TestCase): +class TestModuleCateg(unittest.TestCase): def setUp(self): - self.module_obj = Categ() - self.module_obj.debug = True + self.module = Categ() + self.module.debug = True def test_module(self): item_id = 'tests/2021/01/01/categ.gz' + self.module.obj = Items.Item(item_id) test_categ = ['CreditCards', 'Mail', 'Onion', 'Urls', 'Credential', 'Cve'] - result = self.module_obj.compute(item_id, r_result=True) + result = self.module.compute(None, r_result=True) print(result) self.assertCountEqual(result, test_categ) -class Test_Module_CreditCards(unittest.TestCase): +class TestModuleCreditCards(unittest.TestCase): def setUp(self): - self.module_obj = CreditCards() - self.module_obj.debug = True + self.module = CreditCards() + self.module.debug = True def test_module(self): - item_id = 'tests/2021/01/01/credit_cards.gz 7' + item_id = 'tests/2021/01/01/credit_cards.gz' + self.module.obj = Items.Item(item_id) test_cards = ['341039324930797', # American Express '6011613905509166', # Discover Card '3547151714018657', # Japan Credit Bureau (JCB) @@ -82,27 +84,27 @@ class Test_Module_CreditCards(unittest.TestCase): '4024007132849695', # '4532525919781' # 16-digit VISA, with separators ] - result = self.module_obj.compute(item_id, r_result=True) + result = self.module.compute('7', r_result=True) self.assertCountEqual(result, test_cards) -class Test_Module_DomClassifier(unittest.TestCase): +class TestModuleDomClassifier(unittest.TestCase): def setUp(self): - self.module_obj = DomClassifier() - self.module_obj.debug = True + self.module = DomClassifier() + self.module.debug = True def test_module(self): test_host = 'foo.be' item_id = 'tests/2021/01/01/domain_classifier.gz' - msg = f'{test_host} {item_id}' - result = self.module_obj.compute(msg, r_result=True) + self.module.obj = Items.Item(item_id) + result = self.module.compute(f'{test_host}', r_result=True) self.assertTrue(len(result)) -class Test_Module_Global(unittest.TestCase): +class TestModuleGlobal(unittest.TestCase): def setUp(self): - self.module_obj = Global() - self.module_obj.debug = True + self.module = Global() + self.module.debug = True def test_module(self): # # TODO: delete item @@ -113,25 +115,25 @@ class Test_Module_Global(unittest.TestCase): item_content = b'Lorem ipsum dolor sit amet, consectetur adipiscing elit' item_content_1 = b64encode(gzip.compress(item_content)).decode() item_content_2 = b64encode(gzip.compress(item_content + b' more text ...')).decode() - message = f'{item_id} {item_content_1}' + self.module.obj = Items.Item(item_id) # Test new item - result = self.module_obj.compute(message, r_result=True) + result = self.module.compute(item_content_1, r_result=True) print(f'test new item: {result}') self.assertEqual(result, item_id) # Test duplicate - result = self.module_obj.compute(message, r_result=True) + result = self.module.compute(item_content_1, r_result=True) print(f'test duplicate {result}') self.assertIsNone(result) # Test same id with != content item = Items.Item('tests/2021/01/01/global_831875da824fc86ab5cc0e835755b520.gz') item.delete() - message = f'{item_id} {item_content_2}' - result = self.module_obj.compute(message, r_result=True) + result = self.module.compute(item_content_2, r_result=True) print(f'test same id with != content: {result}') self.assertIn(item_id[:-3], result) + print(result) self.assertNotEqual(result, item_id) # cleanup @@ -139,40 +141,43 @@ class Test_Module_Global(unittest.TestCase): # item.delete() # # TODO: remove from queue -class Test_Module_Keys(unittest.TestCase): +class TestModuleKeys(unittest.TestCase): def setUp(self): - self.module_obj = Keys() - self.module_obj.debug = True + self.module = Keys() + self.module.debug = True def test_module(self): item_id = 'tests/2021/01/01/keys.gz' + self.module.obj = Items.Item(item_id) # # TODO: check results - result = self.module_obj.compute(item_id) + result = self.module.compute(None) -class Test_Module_Onion(unittest.TestCase): +class TestModuleOnion(unittest.TestCase): def setUp(self): - self.module_obj = Onion() - self.module_obj.debug = True + self.module = Onion() + self.module.debug = True def test_module(self): item_id = 'tests/2021/01/01/onion.gz' + self.module.obj = Items.Item(item_id) domain_1 = 'eswpccgr5xyovsahffkehgleqthrasfpfdblwbs4lstd345dwq5qumqd.onion' domain_2 = 'www.facebookcorewwwi.onion' - self.module_obj.compute(f'{item_id} 3') + self.module.compute(f'3') -class Test_Module_Telegram(unittest.TestCase): +class TestModuleTelegram(unittest.TestCase): def setUp(self): - self.module_obj = Telegram() - self.module_obj.debug = True + self.module = Telegram() + self.module.debug = True def test_module(self): item_id = 'tests/2021/01/01/keys.gz' + self.module.obj = Items.Item(item_id) # # TODO: check results - result = self.module_obj.compute(item_id) + result = self.module.compute(None) if __name__ == '__main__':