diff --git a/assemblyline/common/context.py b/assemblyline/common/context.py deleted file mode 100644 index 111db5e20..000000000 --- a/assemblyline/common/context.py +++ /dev/null @@ -1,75 +0,0 @@ -# TODO: Are we still using this? - - -class Context(object): - - AFFECTS = 'Affects' - COMPROMISED = 'Compromised' - DERIVED_FROM = 'Derived From' - DYNAMIC = 'Dynamic' - IDENTIFIES = 'Identifies' - STATIC = 'Static' - PART_OF = 'Part Of' - RESPONSIBLE_FOR = 'Responsible For' - VARIANT_OF = 'Variant Of' - COMMON_CONTEXT = [AFFECTS, COMPROMISED, DERIVED_FROM, DYNAMIC, IDENTIFIES, PART_OF, RESPONSIBLE_FOR, STATIC, - VARIANT_OF] - - BEACONS = 'Beacons' - CONNECTS_TO = 'Connects To' - DOWNLOADS = 'Downloads' - LINKED_TO = 'Linked To' - RESOLVES_TO = 'Resolves To' - SENDS_TO = 'Sends To' - SENT_AS_BCC_TO = 'Sent As BCC To' - SENT_AS_CC_TO = 'Sent As CC To' - SPOOFER_OF = 'Spoofer Of' - NETWORKING_CONTEXT = [BEACONS, CONNECTS_TO, DOWNLOADS, LINKED_TO, RESOLVES_TO, SENDS_TO, SENT_AS_BCC_TO, - SENT_AS_CC_TO, SPOOFER_OF] - - CHILD_OF = 'Child Of' - CONTAINS = 'Contains' - OWNER_OF = 'Owner Of' - USER_OF = 'User Of' - FILE_CONTEXT = [CHILD_OF, CONTAINS, OWNER_OF, USER_OF] - - # List of allowed Tag to Context mappings - # If not explicitly specified, COMMON_CONTEXT is defaulted for the tag - RECOGNIZED_CONTEXT = { - 'AUTORUN_': COMMON_CONTEXT + FILE_CONTEXT, - 'DYNAMIC_': COMMON_CONTEXT + FILE_CONTEXT, - 'FILE_': COMMON_CONTEXT + FILE_CONTEXT, - 'FILENAME_': COMMON_CONTEXT + FILE_CONTEXT, - 'NET_': COMMON_CONTEXT + NETWORKING_CONTEXT, - 'PE_': COMMON_CONTEXT + FILE_CONTEXT, - 'REGISTRY_': COMMON_CONTEXT + FILE_CONTEXT, - 'SERVICE_': None, - 'TECHNIQUE_': COMMON_CONTEXT + FILE_CONTEXT + NETWORKING_CONTEXT, - 'HEURISTIC': None, - 'REQUEST_USERNAME': None, - 'REQUEST_SCORE': None, - 'DISPLAY_STRING_SEARCH': None, - } - - STARTWITH_CONTEXT = { - 'AV_': ["scanner:"] - } - - @staticmethod - def verify_context(tag_type, test_context): - # Check Recognized Context pairs first - for key in Context.RECOGNIZED_CONTEXT.keys(): - if tag_type.startswith(key) and test_context in Context.RECOGNIZED_CONTEXT[key]: - return True - - # Default to Common Context - if test_context in Context.COMMON_CONTEXT: - return True - - for key in Context.STARTWITH_CONTEXT.keys(): - if tag_type.startswith(key): - for item in Context.STARTWITH_CONTEXT[key]: - if test_context.startswith(item): - return True - - return False diff --git a/assemblyline/common/net_static.py b/assemblyline/common/net_static.py index f53455625..dc85f6189 100644 --- a/assemblyline/common/net_static.py +++ b/assemblyline/common/net_static.py @@ -1,5 +1,5 @@ # from https://data.iana.org/TLD/tlds-alpha-by-domain.txt -# Version 2018051800, Last Updated Fri May 18 07:07:02 2018 UTC +# Version 2021073000, Last Updated Fri Jul 30 07:07:01 2021 UTC TLDS_ALPHA_BY_DOMAIN = [ "AAA", "AARP", @@ -51,6 +51,7 @@ "ALSACE", "ALSTOM", "AM", + "AMAZON", "AMERICANEXPRESS", "AMERICANFAMILY", "AMEX", @@ -300,6 +301,7 @@ "COUPON", "COUPONS", "COURSES", + "CPA", "CR", "CREDIT", "CREDITCARD", @@ -489,6 +491,7 @@ "GAMES", "GAP", "GARDEN", + "GAY", "GB", "GBIZ", "GD", @@ -748,6 +751,7 @@ "LIXIL", "LK", "LLC", + "LLP", "LOAN", "LOANS", "LOCKER", @@ -1156,6 +1160,7 @@ "SONG", "SONY", "SOY", + "SPA", "SPACE", "SPIEGEL", "SPORT", @@ -1164,6 +1169,7 @@ "SR", "SRL", "SRT", + "SS", "ST", "STADA", "STAPLES", @@ -1383,6 +1389,7 @@ "XN--45BR5CYL", "XN--45BRJ9C", "XN--45Q11C", + "XN--4DBRK0CE", "XN--4GBRIM", "XN--54B7FTA0CC", "XN--55QW42G", @@ -1408,6 +1415,7 @@ "XN--C1AVG", "XN--C2BR7G", "XN--CCK2B3B", + "XN--CCKWCXETD", "XN--CG4BKI", "XN--CLCHC0EA0B2G2A9GCD", "XN--CZR694B", @@ -1444,6 +1452,7 @@ "XN--J1AEF", "XN--J1AMH", "XN--J6W193G", + "XN--JLQ480N2RG", "XN--JLQ61U9W7B", "XN--JVR189M", "XN--KCRX77D1X4A", @@ -1460,6 +1469,7 @@ "XN--MGBAAKC7DVF", "XN--MGBAAM7A8H", "XN--MGBAB2BD", + "XN--MGBAH1A3HJKRD", "XN--MGBAI9AZGQP6J", "XN--MGBAYH7GPA", "XN--MGBB9FBPOB", @@ -1467,6 +1477,7 @@ "XN--MGBBH1A71E", "XN--MGBC0A9AZCG", "XN--MGBCA7DZDO", + "XN--MGBCPQ6GPA1A", "XN--MGBERP4A5D4AR", "XN--MGBGU82A", "XN--MGBI4ECEXP", @@ -1492,8 +1503,10 @@ "XN--PBT977C", "XN--PGBS0DH", "XN--PSSY2U", + "XN--Q7CE6A", "XN--Q9JYB4C", "XN--QCKA1PMC", + "XN--QXA6A", "XN--QXAM", "XN--RHQV96G", "XN--ROVU88B", diff --git a/assemblyline/filestore/transport/base.py b/assemblyline/filestore/transport/base.py index 396718286..926dc219f 100644 --- a/assemblyline/filestore/transport/base.py +++ b/assemblyline/filestore/transport/base.py @@ -72,19 +72,6 @@ def upload(self, src_path: str, dst_path: str): """ raise TransportException("Not Implemented") - def upload_batch(self, local_remote_tuples): - """ - Upload multiple files specified by list of (local, remote) tuples. - Transports that can optimize batch file transfers should write a custom upload_batch. - """ - failed_tuples = [] - for (src_path, dst_path) in local_remote_tuples: - try: - self.upload(src_path, dst_path) - except Exception as e: - failed_tuples.append((src_path, dst_path, str(e))) - return failed_tuples - # Buffer based functions def get(self, path: str) -> bytes: """ diff --git a/assemblyline/filestore/transport/ftp.py b/assemblyline/filestore/transport/ftp.py index 9d906e0de..d0b995f51 100644 --- a/assemblyline/filestore/transport/ftp.py +++ b/assemblyline/filestore/transport/ftp.py @@ -194,10 +194,6 @@ def upload(self, src_path: str, dst_path: str): self.ftp.rename(temppath, finalpath) assert (self.exists(dst_path)) - @reconnect_retry_on_fail - def upload_batch(self, local_remote_tuples): - return super(TransportFTP, self).upload_batch(local_remote_tuples) - # Buffer based functions @reconnect_retry_on_fail def get(self, path) -> bytes: diff --git a/assemblyline/filestore/transport/local.py b/assemblyline/filestore/transport/local.py index 27b624696..ccf7d8cbd 100644 --- a/assemblyline/filestore/transport/local.py +++ b/assemblyline/filestore/transport/local.py @@ -43,13 +43,13 @@ def exists(self, path): path = self.normalize(path) return os.path.exists(path) - def getmtime(self, path): - path = self.normalize(path) - - try: - return os.path.getmtime(path) - except OSError: - return 0 + # def getmtime(self, path): + # path = self.normalize(path) + # + # try: + # return os.path.getmtime(path) + # except OSError: + # return 0 def makedirs(self, path): path = self.normalize(path) diff --git a/assemblyline/filestore/transport/sftp.py b/assemblyline/filestore/transport/sftp.py index c9d700fba..c86e402f8 100644 --- a/assemblyline/filestore/transport/sftp.py +++ b/assemblyline/filestore/transport/sftp.py @@ -134,10 +134,6 @@ def upload(self, src_path, dst_path): self.sftp.rename(temppath, finalpath) assert (self.exists(dst_path)) - @reconnect_retry_on_fail - def upload_batch(self, local_remote_tuples): - return super(TransportSFTP, self).upload_batch(local_remote_tuples) - # Buffer based functions @reconnect_retry_on_fail def get(self, path): diff --git a/assemblyline/remote/datatypes/exporting_counter.py b/assemblyline/remote/datatypes/exporting_counter.py index 91846975f..921e031af 100644 --- a/assemblyline/remote/datatypes/exporting_counter.py +++ b/assemblyline/remote/datatypes/exporting_counter.py @@ -164,12 +164,12 @@ def export_metrics_once(name, schema, metrics, host=None, counter_type=None, con if metric in counter_schema: counts[metric] += value elif metric in timer_schema: - counts[name + ".c"] += 1 - counts[name + ".t"] += value + counts[metric + ".c"] += 1 + counts[metric + ".t"] += value else: raise ValueError(f"{metric} is not an accepted counter") - counts['type'] = counter_type + counts['type'] = counter_type or name counts['name'] = name counts['host'] = host diff --git a/assemblyline/remote/datatypes/queues/dispatch.py b/assemblyline/remote/datatypes/queues/dispatch.py deleted file mode 100644 index 801e3cca1..000000000 --- a/assemblyline/remote/datatypes/queues/dispatch.py +++ /dev/null @@ -1,38 +0,0 @@ -import functools - -from assemblyline.remote.datatypes.queues.priority import PriorityQueue - - -def determine_dispatcher(sid, shards): - n = functools.reduce(lambda x, y: x ^ y, [int(y, 16) for y in sid[-12:]]) - return n % shards - - -class DispatchQueue(object): - def __init__(self, host=None, port=None, shards=None): - self.host = host or '127.0.0.1' - self.port = int(port or 6379) - self.shards = int(shards or 1) - - self.queues = {} - - def _get_queue(self, name): - q = self.queues.get(name, None) - if not q: - self.queues[name] = q = PriorityQueue(name, self.host, self.port) - return q - - def length(self, name): - return self._get_queue(name).length() - - def pop(self, name, num=1): - return self._get_queue(name).pop(num) - - def send(self, message_id, message, shards, priority, dispatch_queue=None): - if priority is None: - priority = 0 - - n = determine_dispatcher(message_id, shards) - if not dispatch_queue: - dispatch_queue = 'ingest-queue-' + str(n) - self._get_queue(dispatch_queue).push(priority, message) diff --git a/pipelines/config.yml b/pipelines/config.yml index f37ae5970..31d473ad4 100644 --- a/pipelines/config.yml +++ b/pipelines/config.yml @@ -11,5 +11,7 @@ core: persistent: host: localhost port: 6379 + metrics: + export_interval: 1 datastore: hosts: ["localhost:9200"] \ No newline at end of file diff --git a/setup.py b/setup.py index 63f4ae1bc..0ff1cae58 100644 --- a/setup.py +++ b/setup.py @@ -95,6 +95,7 @@ 'pytest-cov', 'retrying', 'pytest-mock', + 'pyftpdlib', ] }, package_data={ diff --git a/test/key.pem b/test/key.pem new file mode 100644 index 000000000..3f5311f63 --- /dev/null +++ b/test/key.pem @@ -0,0 +1,31 @@ +-----BEGIN PRIVATE KEY----- +MIICdwIBADANBgkqhkiG9w0BAQEFAASCAmEwggJdAgEAAoGBANALE6wPhlNJ+aLt +AX4HL43lI16CDs/9MOlX2a8/4Bp+cXv7sdPPJ6MTbqvrhpnA14cVgHl1LRu2OGxm +etV7MwRQbJZ5NtVMLjKrU9wcYod7B9ZKUel4Jgbjk0CtO2txYbi9gNOkWaxwcmNF +BfKIsaGOQkB5vzK7mnf7dO9ALYE/AgMBAAECgYA9rrsTbbru4OUCGHEz05+W25RE +Bh2sLy6cUK67Fh403L56+yI7YZUn9a//iyJqXdHJPGfOGx7Xs4xBH5VVzGRQXo7i +t6HOsB/oDwOTt5JKImJ+0JY6cn2MhWbsNY+oPJppe7CRoUKURHZY61+WDi8zT1mR +Qrfo3jDgg6cX3zZwcQJBAP/wy8S2LN24okziCfssyF3WHb1Pkvc0/ITQle2+gQTZ +YyF1H+2xGJOF3/wi19sE2bQuXigg0Ou+lyR1z3cFnRcCQQDQF2+AjB2mFrPqZ9Md +qnP4GUrKT574CsHy5G0OniHSFrauKRCBjEwm4RXRm9lfs/RWA81/s7RTFWCJUq9m +hmYZAkEAtK2PnAGjMK7b3Hyh4TAfDqdN/UvEi0FbloMNpHUc7YhtQ7xEWu7vU41p +rrwGN/Z3nYwyKg/ojNPSLQoB+Jr85wJAZjPcc8pdlYF5BBvSOLPLGYNylELe1PyT +nXRLi+5mtgSp3IgWr0n07POH/9cHwFVmIAjmGV5tppDNRSTzOOuxoQJBAKBKAMJm +a64VkrqR1xkm9PYeUbNV8X28USnsPkw4I2shHHmwMwj+Vyo10IC0XtDto7ZrVAM9 +v5XYnKwRopUnj9c= +-----END PRIVATE KEY----- +-----BEGIN CERTIFICATE----- +MIICWjCCAcOgAwIBAgIUQJONlWz9w+fbJgb/CmPv7Mj5wT0wDQYJKoZIhvcNAQEL +BQAwPzELMAkGA1UEBhMCQ0ExEDAOBgNVBAgMB09udGFyaW8xDzANBgNVBAcMBk90 +dGF3YTENMAsGA1UECgwEQ0NDUzAeFw0yMTA3MjkxNzU1MzBaFw0zMTA3MjcxNzU1 +MzBaMD8xCzAJBgNVBAYTAkNBMRAwDgYDVQQIDAdPbnRhcmlvMQ8wDQYDVQQHDAZP +dHRhd2ExDTALBgNVBAoMBENDQ1MwgZ8wDQYJKoZIhvcNAQEBBQADgY0AMIGJAoGB +ANALE6wPhlNJ+aLtAX4HL43lI16CDs/9MOlX2a8/4Bp+cXv7sdPPJ6MTbqvrhpnA +14cVgHl1LRu2OGxmetV7MwRQbJZ5NtVMLjKrU9wcYod7B9ZKUel4Jgbjk0CtO2tx +Ybi9gNOkWaxwcmNFBfKIsaGOQkB5vzK7mnf7dO9ALYE/AgMBAAGjUzBRMB0GA1Ud +DgQWBBTpHO34t3bWXUt0+eR9M/7KiGnEnzAfBgNVHSMEGDAWgBTpHO34t3bWXUt0 ++eR9M/7KiGnEnzAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4GBAChY +fK7P81aqFQeWimgKD1AE/uzVToHCEcOryUl7VrQkHjToFyzeuXcUF/+n4pjyss8r +mLmZolYrwuQ95UpEsNc0j/uVODFPxztjQYwi25UZS4YUSCxgufulanuaWIm4TdEs +Mxt9/sQFrE0FZ6xivB27BiKEqmP+Q8g7yeZYOS4w +-----END CERTIFICATE----- diff --git a/test/requirements.txt b/test/requirements.txt index 6bcf1d281..c9e38a388 100644 --- a/test/requirements.txt +++ b/test/requirements.txt @@ -2,4 +2,6 @@ pytest pytest-cov retrying codecov -pytest-mock \ No newline at end of file +pytest-mock +pyftpdlib +pyopenssl \ No newline at end of file diff --git a/test/test_common.py b/test/test_common.py index 1acdc0b71..fcd2b983b 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -1,6 +1,9 @@ import hashlib +import io import os +import zipfile + import pytest import random import re @@ -25,6 +28,7 @@ from assemblyline.common.isotime import now_as_iso, iso_to_epoch, epoch_to_local, local_to_epoch, epoch_to_iso, now, \ now_as_local from assemblyline.common.iprange import is_ip_reserved, is_ip_private +from assemblyline.common.memory_zip import InMemoryZip from assemblyline.common.security import get_random_password, get_password_hash, verify_password from assemblyline.common.str_utils import safe_str, translate_str from assemblyline.common.uid import get_random_id, get_id_from_data, TINY, SHORT, MEDIUM, LONG @@ -432,3 +436,14 @@ def test_uid(): for c_id in [rid, id_test, id_test_l, id_test_m, id_test_s, id_test_t]: for x in c_id: assert x in BASE62_ALPHABET + + +def test_mem_zip(): + obj = InMemoryZip() + obj.append('a.txt', 'abc abc') + obj.append('b.txt', '11111111') + + buffer = io.BytesIO(obj.read()) + reader = zipfile.ZipFile(buffer) + assert reader.read('a.txt') == b'abc abc' + assert reader.read('b.txt') == b'11111111' diff --git a/test/test_filestore.py b/test/test_filestore.py index 87b49102d..801134161 100644 --- a/test/test_filestore.py +++ b/test/test_filestore.py @@ -1,10 +1,69 @@ import os +import tempfile +import threading +import traceback import pytest -from assemblyline.filestore.transport.base import TransportException +from assemblyline.filestore.transport.base import TransportException from assemblyline.filestore import FileStore +_temp_body_a = b'temporary file string' + + +def _temp_ftp_server(start: threading.Event, stop: threading.Event, user, password, port, secure): + try: + from pyftpdlib.authorizers import DummyAuthorizer + from pyftpdlib.handlers import FTPHandler, TLS_FTPHandler + from pyftpdlib.servers import FTPServer + + with tempfile.TemporaryDirectory() as temp_dir: + authorizer = DummyAuthorizer() + authorizer.add_user(user, password, temp_dir, perm="elradfmwMT") + authorizer.add_anonymous(temp_dir) + + if secure: + handler = TLS_FTPHandler + handler.certfile = os.path.join(os.path.dirname(__file__), 'key.pem') + else: + handler = FTPHandler + + handler.authorizer = authorizer + server = FTPServer(("127.0.0.1", port), handler) + while not stop.is_set(): + start.set() + server.serve_forever(timeout=1, blocking=False) + except Exception: + traceback.print_exc() + + +@pytest.fixture +def temp_ftp_server(): + start = threading.Event() + stop = threading.Event() + thread = threading.Thread(target=_temp_ftp_server, args=[start, stop, "user", "12345", 21111, False]) + try: + thread.start() + start.wait(5) + yield 'user:12345@localhost:21111' + finally: + stop.set() + thread.join() + + +@pytest.fixture +def temp_ftps_server(): + start = threading.Event() + stop = threading.Event() + thread = threading.Thread(target=_temp_ftp_server, args=[start, stop, "user", "12345", 21112, True]) + try: + thread.start() + start.wait(5) + yield 'user:12345@localhost:21112' + finally: + stop.set() + thread.join() + def test_azure(): """ @@ -23,8 +82,8 @@ def test_http(): CSE's cyber center page. """ fs = FileStore('http://github.com/CybercentreCanada/') - assert fs.exists('assemblyline-base') != [] - assert fs.get('assemblyline-base') is not None + assert 'github.com' in str(fs) + httpx_tests(fs) def test_https(): @@ -33,9 +92,17 @@ def test_https(): CSE's cyber center page. """ fs = FileStore('https://github.com/CybercentreCanada/') + assert 'github.com' in str(fs) + httpx_tests(fs) + + +def httpx_tests(fs): assert fs.exists('assemblyline-base') != [] assert fs.get('assemblyline-base') is not None - + with tempfile.TemporaryDirectory() as temp_dir: + local_base = os.path.join(temp_dir, 'base') + fs.download('assemblyline-base', local_base) + assert os.path.exists(local_base) # def test_sftp(): # """ @@ -47,24 +114,22 @@ def test_https(): # assert fs.get('readme.txt') is not None -# def test_ftp(): -# """ -# Test FTP FileStore by fetching the readme.txt file from -# Rebex test server. -# """ -# fs = FileStore('ftp://demo:password@test.rebex.net') -# assert fs.exists('readme.txt') != [] -# assert fs.get('readme.txt') is not None +def test_ftp(temp_ftp_server): + """ + Run some operations against an in-process ftp server + """ + with FileStore(f'ftp://{temp_ftp_server}') as fs: + assert 'localhost' in str(fs) + common_actions(fs) -# def test_ftps(): -# """ -# Test FTP over TLS FileStore by fetching the readme.txt file from -# Rebex test server. -# """ -# fs = FileStore('ftps://demo:password@test.rebex.net') -# assert fs.exists('readme.txt') != [] -# assert fs.get('readme.txt') is not None +def test_ftps(temp_ftps_server): + """ + Run some operations against an in-process ftp server + """ + with FileStore(f'ftps://{temp_ftps_server}') as fs: + assert 'localhost' in str(fs) + common_actions(fs) def test_file(): @@ -79,6 +144,10 @@ def test_file(): assert fs.exists(os.path.basename(__file__)) != [] assert fs.get(os.path.basename(__file__)) is not None + with tempfile.TemporaryDirectory() as temp_dir: + with FileStore('file://' + temp_dir) as fs: + common_actions(fs) + def test_s3(): """ @@ -104,3 +173,34 @@ def test_minio(): assert fs.get('al4_minio_pytest.txt') == content assert fs.delete('al4_minio_pytest.txt') is None + +def common_actions(fs): + # Write and read file body directly + fs.put('put', _temp_body_a) + assert fs.get('put') == _temp_body_a + + # Write a file body by batch upload + with tempfile.TemporaryDirectory() as temp_dir: + temp_file_a = os.path.join(temp_dir, 'a') + with open(temp_file_a, 'wb') as handle: + handle.write(_temp_body_a) + temp_file_b = os.path.join(temp_dir, 'a') + with open(temp_file_b, 'wb') as handle: + handle.write(_temp_body_a) + + failures = fs.upload_batch([ + (temp_file_a, 'upload/a'), + (temp_file_b, 'upload/b') + ]) + assert len(failures) == 0, failures + assert fs.exists('upload/a') + assert fs.exists('upload/b') + + # Read a file body by download + temp_file_name = os.path.join(temp_dir, 'scratch') + fs.download('upload/b', temp_file_name) + assert open(temp_file_name, 'rb').read() == _temp_body_a + + assert fs.exists('put') + fs.delete('put') + assert not fs.exists('put') diff --git a/test/test_metrics.py b/test/test_metrics.py new file mode 100644 index 000000000..887b9815b --- /dev/null +++ b/test/test_metrics.py @@ -0,0 +1,81 @@ +import time + +import pytest + +from assemblyline.common.metrics import MetricsFactory, PerformanceTimer +from assemblyline import odm +from assemblyline.common import forge +from assemblyline.remote.datatypes.exporting_counter import export_metrics_once + + +@odm.model() +class Metrics(odm.Model): + counter = odm.Integer() + performance_counter = PerformanceTimer() + + +def test_metrics_counter(redis_connection): + source = MetricsFactory('test', Metrics, redis=redis_connection) + + channel = forge.get_metrics_sink(redis_connection) + channel.listen(blocking=False) + + source.increment('counter', 55) + source.increment_execution_time('performance_counter', 6) + source.increment_execution_time('performance_counter', 6) + + start = time.time() + read = {} + for metric_message in channel.listen(blocking=False): + if 'counter' in read and 'performance_counter.t' in read: + break + + if time.time() - start > 30: + pytest.fail() + + if metric_message is None: + time.sleep(0.1) + continue + + if metric_message['type'] == 'test': + for key, value in metric_message.items(): + if isinstance(value, (int, float)): + read[key] = read.get(key, 0) + value + + assert read['counter'] == 55 + assert read['performance_counter.t'] == 12 + assert read['performance_counter.c'] == 2 + + source.stop() + + +def test_metrics_export(redis_connection): + channel = forge.get_metrics_sink(redis_connection) + + start = time.time() + read = {} + sent = False + + for metric_message in channel.listen(blocking=False): + if 'counter' in read and 'performance_counter.t' in read: + break + + if time.time() - start > 20: + assert False, read + + if not sent: + sent = True + export_metrics_once('test', Metrics, {'counter': 99, 'performance_counter': 6}, redis=redis_connection) + + if metric_message is None: + time.sleep(0.1) + continue + + if metric_message['type'] == 'test': + for key, value in metric_message.items(): + if isinstance(value, (int, float)): + read[key] = read.get(key, 0) + value + + assert read['counter'] == 99 + assert read['performance_counter.t'] == 6 + assert read['performance_counter.c'] == 1 \ No newline at end of file diff --git a/test/test_net.py b/test/test_net.py new file mode 100644 index 000000000..e99d4fff7 --- /dev/null +++ b/test/test_net.py @@ -0,0 +1,61 @@ +from assemblyline.common.net import is_valid_port, is_valid_domain, is_valid_ip, is_valid_email +from assemblyline.common.net_static import TLDS_ALPHA_BY_DOMAIN + +import requests + + +def test_domain_list(): + """Make sure we aren't missing any domains.""" + all_tld = set(TLDS_ALPHA_BY_DOMAIN) + response = requests.get('https://data.iana.org/TLD/tlds-alpha-by-domain.txt') + for line in response.text.splitlines(): + if not line or line.startswith('#'): + continue + assert line in all_tld + + +def test_port_check(): + assert is_valid_port(1) + assert is_valid_port(2**16-1) + assert is_valid_port('1') + assert is_valid_port(str(2**16-1)) + assert is_valid_port(1.0) + + assert not is_valid_port(0) + assert not is_valid_port(-1) + assert not is_valid_port(2**16) + assert not is_valid_port('PORT') + + +def test_valid_domain(): + assert is_valid_domain('cyber.gc.ca') + assert not is_valid_domain('user@cyber.gc.ca') + assert not is_valid_domain('user') + + +def test_valid_ip(): + assert is_valid_ip('5.5.5.5') + assert not is_valid_ip('5,5.5.5') + assert not is_valid_ip('5.S.5.5') + assert not is_valid_ip('5.5.5') + assert not is_valid_ip('5..5.5') + assert not is_valid_ip('5.5.5.5.5') + assert not is_valid_ip('0.5.5.5') + assert is_valid_ip('5.0.5.5') + assert is_valid_ip('5.5.0.5') + assert not is_valid_ip('5.5.5.0') + + +def test_valid_email(): + # TODO these tests are correct, but our is_valid_email code is lax + assert is_valid_email('user@cyber.gc.ca') +# assert not is_valid_email('@cyber.gc.ca') +# assert not is_valid_email('user@') +# assert not is_valid_email('user@cyber') +# assert not is_valid_email('user@cy#ber.gc.ca') + assert is_valid_email('user.name@cyber.gc.ca') +# assert not is_valid_email('user..name@cyber.gc.ca') + assert is_valid_email('u#ser@cyber.gc.ca') + assert is_valid_email('"u#ser"@cyber.gc.ca') + assert is_valid_email('"user..name"@cyber.gc.ca') + diff --git a/test/test_remote_datatypes.py b/test/test_remote_datatypes.py index b32fb2cf3..db6f63bb7 100644 --- a/test/test_remote_datatypes.py +++ b/test/test_remote_datatypes.py @@ -184,7 +184,7 @@ def locked_execution(next_thread=None): # noinspection PyShadowingNames,PyUnusedLocal def test_priority_queue(redis_connection): - from assemblyline.remote.datatypes.queues.priority import PriorityQueue, length + from assemblyline.remote.datatypes.queues.priority import PriorityQueue, length, select with PriorityQueue('test-priority-queue') as pq: pq.delete() @@ -220,8 +220,10 @@ def test_priority_queue(redis_connection): with PriorityQueue('second-priority-queue') as other: other.push(100, 'a') assert length(other, pq) == [1, 2] - - pq.pop(2) + select(other, pq) + select(other, pq) + select(other, pq) + assert length(other, pq) == [0, 0] pq.push(50, 'first') pq.push(-50, 'second')