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

Author: adam-lebon

ctfcli/cli/media.py

@@ -15,16 +15,16 @@ def add(self, path):
 
         api = API()
 
-        new_file = ("file", open(path, mode="rb"))
         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, Union, Optional, List, Any, Tuple
 from urllib.parse import urljoin
 
 from requests import Session
+from requests_toolbelt.multipart.encoder import MultipartEncoder
 
 from ctfcli.core.config import Config
 
@@ -38,14 +40,50 @@ 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", None) is not None:
-            return super(API, self).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()  # type: List[Tuple[str, Any]]
+            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)
+
+            return super(API, self).request(
+                method,
+                url,
+                data=multipart,
+                headers={"Content-Type": multipart.content_type},
+                *args,
+                **kwargs,
+            )
+
+        # If data or files are not key/value pairs, then send the raw values
+        if data is not None or files is not None:
+            return super(API, self).request(
+                method,
+                url,
+                data=data,
+                files=files,
+                *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

ctfcli/core/challenge.py

@@ -351,11 +351,11 @@ 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"))
+        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
@@ -364,7 +364,8 @@ def _create_file(self, local_path: Path):
     def _create_all_files(self):
         new_files = []
         for challenge_file in self["files"]:
-            new_files.append(("file", open(self.challenge_directory / challenge_file, mode="rb")))
+            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"}
 
@@ -374,7 +375,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"]

poetry.lock

@@ -24,6 +24,7 @@ appdirs = "^1.4.4"
 colorama = "^0.4.6"
 fire = "^0.5.0"
 typing-extensions = "^4.7.1"
+requests-toolbelt = "^1.0.0"
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.7.0"

pyproject.toml

@@ -170,4 +170,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)