Merge pull request #15 from CybercentreCanada/AL-1461

use NamedTempFile for content-based submissions
CybercentreCanada · Oct 26, 2021 · c9843ed · c9843ed
2 parents 611e96c + 05baece
commit c9843ed
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/assemblyline_client/v4_client/module/ingest.py b/assemblyline_client/v4_client/module/ingest.py
@@ -1,6 +1,6 @@
 import os
-import tempfile
 from json import dumps
+from tempfile import NamedTemporaryFile
 
 from assemblyline_client.v4_client.common.utils import api_path, api_path_by_module, ClientError
 
@@ -31,12 +31,14 @@ def __call__(self, path=None, content=None, url=None, sha256=None, fname=None, p
 
 If content is provided, the path is used as metadata only.
 """
+        temp_file = None
         if content:
-            fd, path = tempfile.mkstemp()
-            with os.fdopen(fd, 'wb') as fh:
-                if isinstance(content, str):
-                    content = content.encode()
-                fh.write(content)
+            temp_file = NamedTemporaryFile(mode="w+b", delete=False)
+            if isinstance(content, str):
+                content = content.encode()
+            temp_file.write(content)
+            temp_file.seek(0)
+            path = temp_file.name
 
         files = {}
         if path:
@@ -48,6 +50,8 @@ def __call__(self, path=None, content=None, url=None, sha256=None, fname=None, p
             request = {
                 'name': fname or os.path.basename(path)
             }
+            if temp_file:
+                temp_file.close()
         elif url:
             request = {
                 'url': url,

diff --git a/assemblyline_client/v4_client/module/submit.py b/assemblyline_client/v4_client/module/submit.py
@@ -1,6 +1,6 @@
 import os
-import tempfile
 from json import dumps
+from tempfile import NamedTemporaryFile
 
 from assemblyline_client.v4_client.common.utils import api_path, api_path_by_module, get_function_kwargs, ClientError
 
@@ -26,12 +26,14 @@ def __call__(self, path=None, content=None, url=None, sha256=None, fname=None, p
 
 If content is provided, the path is used as metadata only.
 """
+        temp_file = None
         if content:
-            fd, path = tempfile.mkstemp()
-            with os.fdopen(fd, 'wb') as fh:
-                if isinstance(content, str):
-                    content = content.encode()
-                fh.write(content)
+            temp_file = NamedTemporaryFile(mode="w+b", delete=False)
+            if isinstance(content, str):
+                content = content.encode()
+            temp_file.write(content)
+            temp_file.seek(0)
+            path = temp_file.name
 
         files = {}
         if path:
@@ -43,6 +45,8 @@ def __call__(self, path=None, content=None, url=None, sha256=None, fname=None, p
             request = {
                 'name': fname or os.path.basename(path)
             }
+            if temp_file:
+                temp_file.close()
         elif url:
             request = {
                 'url': url,