ail-framework/tests/test_modules.py

189 lines
5.5 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import unittest
import gzip
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
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
2021-05-19 12:54:34 +00:00
# project packages
import lib.objects.Items as Items
#### COPY SAMPLES ####
config_loader = ConfigLoader()
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)
2024-02-27 10:03:11 +00:00
#### ---- ####
class TestModuleApiKey(unittest.TestCase):
2021-05-19 12:54:34 +00:00
def setUp(self):
self.module = ApiKey()
self.module.debug = True
2021-05-19 12:54:34 +00:00
def test_module(self):
item_id = 'tests/2021/01/01/api_keys.gz'
self.module.obj = Items.Item(item_id)
2021-05-19 12:54:34 +00:00
google_api_key = 'AIza00000000000000000000000_example-KEY'
aws_access_key = 'AKIAIOSFODNN7EXAMPLE'
aws_secret_key = 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY'
matches = self.module.compute('3', r_result=True)
2023-05-23 11:38:33 +00:00
self.assertCountEqual(matches[0], {google_api_key})
self.assertCountEqual(matches[1], {aws_access_key})
self.assertCountEqual(matches[2], {aws_secret_key})
2021-05-19 12:54:34 +00:00
2024-02-27 10:03:11 +00:00
class TestModuleCateg(unittest.TestCase):
2021-05-19 14:57:20 +00:00
def setUp(self):
self.module = Categ()
self.module.debug = True
2021-05-19 14:57:20 +00:00
def test_module(self):
item_id = 'tests/2021/01/01/categ.gz'
self.module.obj = Items.Item(item_id)
2022-05-30 07:03:27 +00:00
test_categ = ['CreditCards', 'Mail', 'Onion', 'Urls', 'Credential', 'Cve']
2021-05-19 14:57:20 +00:00
result = self.module.compute(None, r_result=True)
2021-05-19 14:57:20 +00:00
self.assertCountEqual(result, test_categ)
2024-02-27 10:03:11 +00:00
class TestModuleCreditCards(unittest.TestCase):
def setUp(self):
self.module = CreditCards()
self.module.debug = True
def test_module(self):
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)
'5492981206527330', # 16 digits MasterCard
'4024007132849695', # '4532525919781' # 16-digit VISA, with separators
2024-02-27 10:03:11 +00:00
]
result = self.module.compute('7', r_result=True)
self.assertCountEqual(result, test_cards)
2024-02-27 10:03:11 +00:00
class TestModuleDomClassifier(unittest.TestCase):
def setUp(self):
self.module = DomClassifier()
self.module.debug = True
def test_module(self):
2022-05-30 07:03:27 +00:00
test_host = 'foo.be'
item_id = 'tests/2021/01/01/domain_classifier.gz'
self.module.obj = Items.Item(item_id)
result = self.module.compute(f'{test_host}', r_result=True)
self.assertTrue(len(result))
2024-02-27 10:03:11 +00:00
class TestModuleGlobal(unittest.TestCase):
def setUp(self):
self.module = Global()
self.module.debug = True
def test_module(self):
# # TODO: delete item
item_id = 'tests/2021/01/01/global.gz'
item = Items.Item(item_id)
item.delete()
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()
self.module.obj = Items.Item(item_id)
# Test new item
result = self.module.compute(item_content_1, r_result=True)
self.assertEqual(result, item_id)
# Test duplicate
result = self.module.compute(item_content_1, r_result=True)
self.assertIsNone(result)
# Test same id with != content
item = Items.Item('tests/2021/01/01/global_831875da824fc86ab5cc0e835755b520.gz')
item.delete()
result = self.module.compute(item_content_2, r_result=True)
2021-05-28 15:23:51 +00:00
self.assertIn(item_id[:-3], result)
self.assertNotEqual(result, item_id)
# cleanup
# item = Items.Item(result)
# item.delete()
# # TODO: remove from queue
2024-02-27 10:03:11 +00:00
class TestModuleKeys(unittest.TestCase):
2021-05-28 15:23:51 +00:00
def setUp(self):
self.module = Keys()
self.module.debug = True
2021-05-28 15:23:51 +00:00
def test_module(self):
item_id = 'tests/2021/01/01/keys.gz'
self.module.obj = Items.Item(item_id)
2021-05-28 15:23:51 +00:00
# # TODO: check results
2024-02-27 10:03:11 +00:00
self.module.compute(None)
2021-05-28 15:23:51 +00:00
class TestModuleOnion(unittest.TestCase):
def setUp(self):
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)
2024-02-27 10:03:11 +00:00
# domain_1 = 'eswpccgr5xyovsahffkehgleqthrasfpfdblwbs4lstd345dwq5qumqd.onion'
# domain_2 = 'www.facebookcorewwwi.onion'
self.module.compute(f'3')
2021-05-19 12:54:34 +00:00
2024-02-27 10:03:11 +00:00
class TestModuleTelegram(unittest.TestCase):
def setUp(self):
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
2024-02-27 10:03:11 +00:00
self.module.compute(None)
2021-05-19 12:54:34 +00:00
if __name__ == '__main__':
unittest.main()