Merge pull request #53 from CybercentreCanada/AL-2646

Modify APIs to support passing submission profile names

Author: cccs-douglass

assemblyline_client/v4_client/module/ingest.py

@@ -11,7 +11,7 @@ def __init__(self, connection):
         self._connection = connection
 
     def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fname=None, params=None, metadata=None,
-                 alert=False, nq=None, nt=None, ingest_type='AL_CLIENT'):
+                 alert=False, nq=None, nt=None, ingest_type='AL_CLIENT', submission_profile=None):
         """\
 Submit a file to the ingestion queue.
 
@@ -30,6 +30,7 @@ def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fnam
 nt         : Notification threshold. (int)
 params     : Additional submission parameters. (dict)
 ingest_type: Ingestion type, one word to describe how the data is ingested. Default: AL_CLIENT (string)
+submission_profile    : Submission profile name
 
 If content is provided, the path is used as metadata only.
 """
@@ -87,7 +88,8 @@ def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fnam
                 request['notification_threshold'] = int(nt)
             if params:
                 request['params'] = params
-
+            if submission_profile:
+                request['submission_profile'] = submission_profile
             if files:
                 data = {'json': dumps(request)}
                 headers = {'content-type': None}

assemblyline_client/v4_client/module/submit.py

@@ -10,7 +10,7 @@ class Submit(object):
     def __init__(self, connection):
         self._connection = connection
 
-    def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fname=None, params=None, metadata=None):
+    def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fname=None, params=None, metadata=None, submission_profile=None):
         """\
 Submit a file to be dispatched.
 
@@ -22,9 +22,10 @@ def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fnam
 url     : Url to scan (string)
 
 Optional
-fname   : Name of the file to scan
-metadata   : Metadata to include with submission. (dict)
-params  : Additional submission parameters. (dict)
+fname       : Name of the file to scan
+metadata    : Metadata to include with submission. (dict)
+params      : Additional submission parameters. (dict)
+submission_profile     : Submission profile name
 
 If content is provided, the path is used as metadata only.
 """
@@ -73,6 +74,9 @@ def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fnam
             if metadata:
                 request['metadata'] = metadata
 
+            if submission_profile:
+                request['submission_profile'] = submission_profile
+
             if files:
                 data = {'json': dumps(request)}
                 headers = {'content-type': None}

test/conftest.py

@@ -1,6 +1,6 @@
 import os
-import pytest
 
+import pytest
 
 UI_HOST = os.getenv('UI_HOST', "https://localhost:443")
 
@@ -24,13 +24,13 @@ def skip_or_fail(message):
 
 
 try:
-    from assemblyline.common.security import get_random_password, get_password_hash
-    from assemblyline.datastore.store import ESStore
-    from assemblyline.datastore.helper import AssemblylineDatastore
-    from assemblyline.common.uid import get_random_id
-
     from assemblyline_client import get_client
+
     from assemblyline.common import forge
+    from assemblyline.common.security import get_password_hash, get_random_password
+    from assemblyline.common.uid import get_random_id
+    from assemblyline.datastore.helper import AssemblylineDatastore
+    from assemblyline.datastore.store import ESStore
     from assemblyline.odm import random_data
 
     config = forge.get_config()
@@ -80,15 +80,8 @@ def datastore(datastore_connection, filestore):
             random_data.wipe_workflows(ds)
 
     @pytest.fixture(scope="module")
-    def client(datastore):
-        user = datastore.user.get('admin')
-        random_pass = get_random_password(length=48)
-        key_name = "key_%s" % get_random_id().lower()
-        user.apikeys[key_name] = {"password": get_password_hash(random_pass), "acl": ["R", "W", "E"]}
-        datastore.user.save('admin', user)
-        api_key = "%s:%s" % (key_name, random_pass)
-
-        c = get_client(UI_HOST, apikey=('admin', api_key), verify=False, retries=1)
+    def client():
+        c = get_client(UI_HOST, auth=('admin', 'admin'), verify=False, retries=1)
         return c
 
 except (ImportError, ModuleNotFoundError):

test/test_user.py

@@ -1,14 +1,14 @@
-import pytest
 import random
 
+import pytest
 from assemblyline_client.v4_client.common.utils import ClientError
 
 try:
-    from assemblyline.odm.models.user_favorites import Favorite
+    from utils import random_id_from_collection
+
     from assemblyline.odm.models.user import User
+    from assemblyline.odm.models.user_favorites import Favorite
     from assemblyline.odm.random_data import random_model_obj
-
-    from utils import random_id_from_collection
 except ImportError:
     import sys
     if sys.version_info < (3, 0):
@@ -143,14 +143,15 @@ def test_settings(datastore, client):
     assert {'classification', 'download_encoding', 'priority', 'ttl'}.issubset(set(settings.keys()))
 
     # Test updating settings
-    new_desc = 'NEW_DEFAULT_DESCRIPTION'
-    settings['description'] = new_desc
+    new_password = "zippy"
+    settings['default_zip_password'] = new_password
 
     res = client.user.settings.update(user_id, settings)
     assert res['success']
 
     # Validate changes were applied
-    assert new_desc == datastore.user_settings.get(user_id, as_obj=False)['description']
+    datastore.user_settings.commit()
+    assert new_password == datastore.user_settings.get(user_id, as_obj=False)['default_zip_password']
 
 
 def test_submission_params(datastore, client):