Generate filehandler via BytesIO if given content-only (#22)

Credits to @malvidin
CybercentreCanada · Mar 23, 2022 · 25510c4 · 25510c4
1 parent 587be50
commit 25510c4
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 14 deletions.
diff --git a/assemblyline_client/v4_client/common/submit_utils.py b/assemblyline_client/v4_client/common/submit_utils.py
@@ -1,7 +1,11 @@
 import json
+import hashlib
 import pprint
 import sys
 
+from io import BytesIO
+from typing import Union
+
 SRV_BUSY_ID = "20"
 SRV_DOWN_ID = "21"
 MAX_RETRY_ID = "12"
@@ -32,6 +36,14 @@
     RD = u"\ua714".encode("utf-8")
 
 
+def get_file_handler(content: Union[str, bytes], fname=None) -> BytesIO:
+    if isinstance(content, str):
+        content = content.encode()
+    fh = BytesIO(content)
+    fh.name = fname or hashlib.sha256(content).hexdigest()
+    return fh
+
+
 def al_result_to_text(r, show_errors=True, verbose_error=False):
     lines = ["", ":: Submission Detail %s::" % {True: "", False: "[Errors hidden]"}[show_errors],
              "\t%-36s %s" % ("state:", r["state"]), ""]

diff --git a/assemblyline_client/v4_client/module/ingest.py b/assemblyline_client/v4_client/module/ingest.py
@@ -1,9 +1,9 @@
 import os
-import tempfile
 
 from json import dumps
 
 from assemblyline_client.v4_client.common.utils import api_path, api_path_by_module, ClientError
+from assemblyline_client.v4_client.common.submit_utils import get_file_handler
 
 
 class Ingest(object):
@@ -36,12 +36,7 @@ def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fnam
         rmpath = None
         try:
             if content:
-                fd, path = tempfile.mkstemp()
-                rmpath = path
-                with os.fdopen(fd, 'wb') as content_fh:
-                    if isinstance(content, str):
-                        content = content.encode()
-                    content_fh.write(content)
+                fh = get_file_handler(content, fname)
 
             files = {}
             if fh:

diff --git a/assemblyline_client/v4_client/module/submit.py b/assemblyline_client/v4_client/module/submit.py
@@ -1,9 +1,9 @@
 import os
-import tempfile
 
 from json import dumps
 
 from assemblyline_client.v4_client.common.utils import api_path, api_path_by_module, get_function_kwargs, ClientError
+from assemblyline_client.v4_client.common.submit_utils import get_file_handler
 
 
 class Submit(object):
@@ -31,12 +31,7 @@ def __call__(self, fh=None, path=None, content=None, url=None, sha256=None, fnam
         rmpath = None
         try:
             if content:
-                fd, path = tempfile.mkstemp()
-                rmpath = path
-                with os.fdopen(fd, 'wb') as content_fh:
-                    if isinstance(content, str):
-                        content = content.encode()
-                    content_fh.write(content)
+                fh = get_file_handler(content, fname)
 
             files = {}
             if fh: