Stream POST request in order to handle large files

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
 from urllib.parse import urljoin
 
 from requests import Session
+from requests_toolbelt.multipart.encoder import MultipartEncoder
 
 from ctfcli.core.config import Config
 
@@ -38,20 +40,53 @@ 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()
+            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,
+            )
 
         # 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", None) is None:
-            kwargs["headers"] = {}
+        if data is None and files is None:
+            if kwargs.get("headers", None) is None:
+                kwargs["headers"] = {}
+            kwargs["headers"]["Content-Type"] = "application/json"
 
-        kwargs["headers"]["Content-Type"] = "application/json"
-        return super(API, self).request(method, url, *args, **kwargs)
+        return super(API, self).request(
+            method,
+            url,
+            data=data,
+            files=files,
+            *args,
+            **kwargs,
+        )

ctfcli/core/challenge.py

@@ -352,11 +352,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
@@ -365,7 +365,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"}
 
@@ -375,7 +376,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"]

pyproject.toml

@@ -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"

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)