Stream POST request in order to handle large files (#161)

* Add MultipartEncoder to support request streaming

The Multipart encoder helps requests to upload large files without the need to read the entire file in memory

* Remove unused typing

* Remove duplicate call

* Update _create_solution to use new file upload strategy

* Only override content type instead of all headers

---------

Co-authored-by: Kevin Chung <kchung@ctfd.io>

Author: maelle-le-bon

ctfcli/cli/media.py

@@ -15,16 +15,16 @@ def add(self, path):
 
         api = API()
 
-        new_file = ("file", open(path, mode="rb"))  # noqa: SIM115
         filename = os.path.basename(path)
+        new_file = (filename, open(path, mode="rb"))
         location = f"media/{filename}"
         file_payload = {
             "type": "page",
             "location": location,
         }
 
         # Specifically use data= here to send multipart/form-data
-        r = api.post("/api/v1/files", files=[new_file], data=file_payload)
+        r = api.post("/api/v1/files", files={"file": new_file}, data=file_payload)
         r.raise_for_status()
         resp = r.json()
         server_location = resp["data"][0]["location"]

ctfcli/core/api.py

@@ -1,6 +1,8 @@
+from typing import Mapping
 from urllib.parse import urljoin
 
 from requests import Session
+from requests_toolbelt.multipart.encoder import MultipartEncoder
 
 from ctfcli.core.config import Config
 from ctfcli.core.exceptions import MissingAPIKey, MissingInstanceURL
@@ -46,20 +48,56 @@ def __init__(self):
         if "cookies" in config:
             self.cookies.update(dict(config["cookies"]))
 
-    def request(self, method, url, *args, **kwargs):
+    def request(self, method, url, data=None, files=None, *args, **kwargs):
         # Strip out the preceding / so that urljoin creates the right url
         # considering the appended / on the prefix_url
         url = urljoin(self.prefix_url, url.lstrip("/"))
 
-        # if data= is present, do not modify the content-type
-        if kwargs.get("data") is not None:
-            return super().request(method, url, *args, **kwargs)
+        # If data or files are any kind of key/value iterable
+        # then encode the body as form-data
+        if isinstance(data, (list, tuple, Mapping)) or isinstance(files, (list, tuple, Mapping)):
+            # In order to use the MultipartEncoder, we need to convert data and files to the following structure :
+            # A list of tuple containing the key and the values : List[Tuple[str, str]]
+            # For files, the structure can be List[Tuple[str, Tuple[str, str, Optional[str]]]]
+            # Example: [ ('file', ('doc.pdf', open('doc.pdf'), 'text/plain') ) ]
+
+            fields = list()
+            if isinstance(data, dict):
+                # int are not allowed as value in MultipartEncoder
+                fields = list(map(lambda v: (v[0], str(v[1]) if isinstance(v[1], int) else v[1]), data.items()))
+
+            if files is not None:
+                if isinstance(files, dict):
+                    files = list(files.items())
+                fields.extend(files)  # type: ignore
+
+            multipart = MultipartEncoder(fields)
+            headers = kwargs.pop("headers", {}) or {}
+            headers = dict(headers)
+            headers["Content-Type"] = multipart.content_type
+
+            return super(API, self).request(
+                method,
+                url,
+                data=multipart,
+                headers=headers,
+                *args,
+                **kwargs,
+            )
 
         # otherwise set the content-type to application/json for all API requests
         # modify the headers here instead of using self.headers because we don't want to
         # override the multipart/form-data case above
-        if kwargs.get("headers") is None:
-            kwargs["headers"] = {}
-
-        kwargs["headers"]["Content-Type"] = "application/json"
-        return super().request(method, url, *args, **kwargs)
+        if data is None and files is None:
+            if kwargs.get("headers", None) is None:
+                kwargs["headers"] = {}
+            kwargs["headers"]["Content-Type"] = "application/json"
+
+        return super(API, self).request(
+            method,
+            url,
+            data=data,
+            files=files,
+            *args,
+            **kwargs,
+        )

ctfcli/core/challenge.py

@@ -389,22 +389,21 @@ def _delete_file(self, remote_location: str):
                 r.raise_for_status()
 
     def _create_file(self, local_path: Path):
-        new_file = ("file", open(local_path, mode="rb"))  # noqa: SIM115
+        new_file = (local_path.name, open(local_path, mode="rb"))
         file_payload = {"challenge_id": self.challenge_id, "type": "challenge"}
 
         # Specifically use data= here to send multipart/form-data
-        r = self.api.post("/api/v1/files", files=[new_file], data=file_payload)
+        r = self.api.post("/api/v1/files", files={"file": new_file}, data=file_payload)
         r.raise_for_status()
 
         # Close the file handle
         new_file[1].close()
 
     def _create_all_files(self):
         new_files = []
-
-        files = self.get("files") or []
-        for challenge_file in files:
-            new_files.append(("file", open(self.challenge_directory / challenge_file, mode="rb")))  # noqa: SIM115
+        for challenge_file in self["files"]:
+            file_path = self.challenge_directory / challenge_file
+            new_files.append(("file", (file_path.name, file_path.open("rb"))))
 
         files_payload = {"challenge_id": self.challenge_id, "type": "challenge"}
 
@@ -414,7 +413,7 @@ def _create_all_files(self):
 
         # Close the file handles
         for file_payload in new_files:
-            file_payload[1].close()
+            file_payload[1][1].close()
 
     def _delete_existing_hints(self):
         remote_hints = self.api.get("/api/v1/hints").json()["data"]
@@ -585,14 +584,15 @@ def _create_solution(self):
             snippet_includes = re.findall(r'(--8<--\s+["\']([^"\']+)["\'])', content)
 
             for mdx, alt, path in markdown_images:
-                new_file = ("file", open(solution_path.parent / path, mode="rb"))
+                local_path = solution_path.parent / path
+                new_file = (local_path.name, open(solution_path.parent / path, mode="rb"))
                 file_payload = {
                     "type": "solution",
                     "solution_id": solution_id,
                 }
 
                 # Specifically use data= here to send multipart/form-data
-                r = self.api.post("/api/v1/files", files=[new_file], data=file_payload)
+                r = self.api.post("/api/v1/files", files={"file": new_file}, data=file_payload)
                 r.raise_for_status()
                 resp = r.json()
                 server_location = resp["data"][0]["location"]

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "fire>=0.7.0,<0.8",
     "typing-extensions>=4.7.1,<5",
     "python-slugify>=8.0.4,<9",
+    "requests-toolbelt==1.0.0",
 ]
 
 [project.scripts]
@@ -49,4 +50,3 @@ exclude = ["build/**", "ctfcli/templates/**"]
 
 [tool.ruff.format]
 exclude = ["build/**", "ctfcli/templates/**"]
-

tests/core/test_api.py

@@ -46,8 +46,20 @@ def test_api_object_request_strips_preceding_slash_from_url_path(self, mock_requ
 
         mock_request.assert_has_calls(
             [
-                call("GET", "https://example.com/test/path", headers={"Content-Type": "application/json"}),
-                call("GET", "https://example.com/test/path", headers={"Content-Type": "application/json"}),
+                call(
+                    "GET",
+                    "https://example.com/test/path",
+                    headers={"Content-Type": "application/json"},
+                    data=None,
+                    files=None,
+                ),
+                call(
+                    "GET",
+                    "https://example.com/test/path",
+                    headers={"Content-Type": "application/json"},
+                    data=None,
+                    files=None,
+                ),
             ]
         )
 
@@ -60,7 +72,7 @@ def test_api_object_request_assigns_prefix_url(self, mock_request: MagicMock, *a
         api = API()
         api.request("GET", "path")
         mock_request.assert_called_once_with(
-            "GET", "https://example.com/test/path", headers={"Content-Type": "application/json"}
+            "GET", "https://example.com/test/path", headers={"Content-Type": "application/json"}, data=None, files=None
         )
 
     def test_api_object_assigns_ssl_verify(self, *args, **kwargs):
@@ -170,4 +182,4 @@ def test_api_object_assigns_cookies(self, *args, **kwargs):
     def test_request_does_not_override_form_data_content_type(self, mock_request: MagicMock, *args, **kwargs):
         api = API()
         api.request("GET", "/test", data="some-file")
-        mock_request.assert_called_once_with("GET", "https://example.com/test", data="some-file")
+        mock_request.assert_called_once_with("GET", "https://example.com/test", data="some-file", files=None)