Repository: CybercentreCanada/assemblyline-service-server (language: python)

Issue:
UTF-8 File Download



Discussion / hints:

Fix patch:
diff --git a/assemblyline_service_server/api/v1/badlist.py b/assemblyline_service_server/api/v1/badlist.py
new file mode 100644
index 0000000..5b68160
--- /dev/null
+++ b/assemblyline_service_server/api/v1/badlist.py
@@ -0,0 +1,125 @@
+from flask import request
+
+from assemblyline_service_server.api.base import api_login, make_subapi_blueprint
+from assemblyline_service_server.config import BADLIST_CLIENT
+from assemblyline_service_server.helper.response import make_api_response
+
+SUB_API = 'badlist'
+badlist_api = make_subapi_blueprint(SUB_API, api_version=1)
+badlist_api._doc = "Query badlisted hashes"
+
+
+@badlist_api.route("/<qhash>/", methods=["GET"])
+@api_login()
+def exists(qhash, **_):
+    """
+    Check if a file exists in the badlist.
+
+    Variables:
+    qhash       => Hash to check
+
+    Arguments:
+    None
+
+    Data Block:
+    None
+
+    API call example:
+    GET /api/v1/badlist/123456...654321/
+
+    Result example:
+    <Badlisting object>
+    """
+    badlist = BADLIST_CLIENT.exists(qhash)
+    if badlist:
+        return make_api_response(badlist)
+    return make_api_response(None, "The hash was not found in the badlist.", 404)
+
+
+@badlist_api.route("/ssdeep/", methods=["POST"])
+@api_login()
+def similar_ssdeep(**_):
+    """
+    Check if a file with a similar SSDeep exists.
+
+    Variables:
+    None
+
+    Arguments:
+    None
+
+    Data Block:
+    {
+        ssdeep : value    => Hash to check
+    }
+
+    API call example:
+    GET /api/v1/badlist/ssdeep/
+
+    Result example:
+    <Badlisting object>
+    """
+    badlist = BADLIST_CLIENT.find_similar_ssdeep(request.json.get('ssdeep', None))
+    if badlist:
+        return make_api_response(badlist)
+    return make_api_response(None, "The hash was not found in the badlist.", 404)
+
+
+@badlist_api.route("/tlsh/", methods=["POST"])
+@api_login()
+def similar_tlsh(**_):
+    """
+    Check if a file with a similar TLSH exists.
+
+    Variables:
+    None
+
+    Arguments:
+    None
+
+    Data Block:
+    {
+        tlsh : value    => Hash to check
+    }
+
+    API call example:
+    GET /api/v1/badlist/tlsh/
+
+    Result example:
+    <Badlisting object>
+    """
+    badlist = BADLIST_CLIENT.find_similar_tlsh(request.json.get('tlsh', None))
+    if badlist:
+        return make_api_response(badlist)
+    return make_api_response(None, "The hash was not found in the badlist.", 404)
+
+
+@badlist_api.route("/tags/", methods=["POST"])
+@api_login()
+def tags_exists(**_):
+    """
+    Check if the provided tags exists in the badlist
+
+    Variables:
+    None
+
+    Arguments:
+    None
+
+    Data Block:
+    { # Dictionary of types -> values to check if exists
+        "file.synamic.domain": [...],
+        "network.static.ip": [...]
+    }
+
+    API call example:
+    GET /api/v1/badlist/tags/
+
+    Result example:
+    [ # List of existing objecs
+        <badlisting object>,
+        <Badlisting object>
+    ]
+    """
+    data = request.json
+    return make_api_response(BADLIST_CLIENT.exists_tags(data))
diff --git a/assemblyline_service_server/app.py b/assemblyline_service_server/app.py
index 3ff68e9..6053aad 100644
--- a/assemblyline_service_server/app.py
+++ b/assemblyline_service_server/app.py
@@ -6,6 +6,7 @@
 from os import environ, path
 
 from assemblyline.common import forge, log as al_log
+from assemblyline_service_server.api.v1.badlist import badlist_api
 from assemblyline_service_server.api.v1.file import file_api
 from assemblyline_service_server.api.v1.service import service_api
 from assemblyline_service_server.api.v1.task import task_api
@@ -32,6 +33,7 @@
 app = Flask('svc-server')
 app.config['SECRET_KEY'] = config.ui.secret_key
 
+app.register_blueprint(badlist_api)
 app.register_blueprint(healthz)
 app.register_blueprint(file_api)
 app.register_blueprint(service_api)
diff --git a/assemblyline_service_server/config.py b/assemblyline_service_server/config.py
index 2d88d1c..ab21340 100644
--- a/assemblyline_service_server/config.py
+++ b/assemblyline_service_server/config.py
@@ -7,6 +7,7 @@
 from assemblyline.common.version import BUILD_MINOR, FRAMEWORK_VERSION, SYSTEM_VERSION
 from assemblyline.remote.datatypes.counters import Counters
 from assemblyline.remote.datatypes import get_client
+from assemblyline_core.badlist_client import BadlistClient
 from assemblyline_core.safelist_client import SafelistClient
 from assemblyline_core.tasking_client import TaskingClient
 
@@ -59,5 +60,6 @@
 LOCK = threading.Lock()
 TASKING_CLIENT = TaskingClient(datastore=STORAGE, filestore=FILESTORE, redis=redis, redis_persist=redis_persist)
 SAFELIST_CLIENT = SafelistClient(datastore=STORAGE)
+BADLIST_CLIENT = BadlistClient(datastore=STORAGE)
 # End global
 #################################################################


Test patch:
diff --git a/test/test_srv_badlist.py b/test/test_srv_badlist.py
new file mode 100644
index 0000000..3e843ac
--- /dev/null
+++ b/test/test_srv_badlist.py
@@ -0,0 +1,76 @@
+
+import pytest
+
+from unittest.mock import MagicMock, patch
+
+from assemblyline.odm import randomizer
+from assemblyline.odm.models.badlist import Badlist
+from assemblyline_service_server import app
+from assemblyline_service_server.config import AUTH_KEY
+
+headers = {
+    'Container-Id': randomizer.get_random_hash(12),
+    'X-APIKey': AUTH_KEY,
+    'Service-Name': 'Badlist',
+    'Service-Version': randomizer.get_random_service_version(),
+    'Service-Tool-Version': randomizer.get_random_hash(64),
+    'Timeout': 1,
+    'X-Forwarded-For': '127.0.0.1',
+}
+
+
+@pytest.fixture(scope='function')
+def storage():
+    ds = MagicMock()
+    with patch('assemblyline_service_server.config.BADLIST_CLIENT.datastore', ds):
+        yield ds
+
+
+@pytest.fixture()
+def client():
+    client = app.app.test_client()
+    yield client
+
+
+# noinspection PyUnusedLocal
+def test_badlist_exist(client, storage):
+    valid_hash = randomizer.get_random_hash(64)
+    valid_resp = randomizer.random_model_obj(Badlist, as_json=True)
+    valid_resp['hashes']['sha256'] = valid_hash
+    storage.badlist.get_if_exists.return_value = valid_resp
+
+    resp = client.get(f'/api/v1/badlist/{valid_hash}/', headers=headers)
+    assert resp.status_code == 200
+    assert resp.json['api_response'] == valid_resp
+
+
+# noinspection PyUnusedLocal
+def test_badlist_missing(client, storage):
+    invalid_hash = randomizer.get_random_hash(64)
+    storage.badlist.get_if_exists.return_value = None
+
+    resp = client.get(f'/api/v1/badlist/{invalid_hash}/', headers=headers)
+    assert resp.status_code == 404
+    assert resp.json['api_response'] is None
+
+
+# noinspection PyUnusedLocal
+def test_get_full_badlist(client, storage):
+    storage.badlist.stream_search.return_value = []
+
+    resp = client.get('/api/v1/badlist/', headers=headers)
+    assert resp.status_code == 200
+    assert isinstance(resp.json['api_response'], dict)
+
+
+# noinspection PyUnusedLocal
+def test_get_full_badlist_specific(client, storage):
+    storage.badlist.stream_search.return_value = []
+
+    tag_type = "network.dynamic.domain"
+    resp = client.get(f'/api/v1/badlist/?tag_types={tag_type}', headers=headers)
+    assert resp.status_code == 200
+    assert isinstance(resp.json['api_response'], dict)
+
+    for k in resp.json['api_response']:
+        assert k == tag_type