From ce0186bc014a74ae4ff247ff1d107973ee086470 Mon Sep 17 00:00:00 2001 From: d3liri0us Date: Fri, 17 Jul 2026 14:50:01 +0530 Subject: [PATCH] add media for challenges --- ctfcli/cli/media.py | 31 ++----------------- ctfcli/core/challenge.py | 58 +++++++++++++++++++++++++----------- ctfcli/core/media.py | 42 ++++++++++++++++++++++++-- tests/core/test_challenge.py | 49 +++++++++++++++++++++++++++++- 4 files changed, 131 insertions(+), 49 deletions(-) diff --git a/ctfcli/cli/media.py b/ctfcli/cli/media.py index e142f0b..cc982ab 100644 --- a/ctfcli/cli/media.py +++ b/ctfcli/cli/media.py @@ -1,41 +1,14 @@ -import os - import click from ctfcli.core.api import API from ctfcli.core.config import Config +from ctfcli.core.media import Media class MediaCommand: def add(self, path): """Add local media file to config file and remote instance""" - config = Config() - if config.config.has_section("media") is False: - config.config.add_section("media") - - api = API() - - 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={"file": new_file}, data=file_payload) - r.raise_for_status() - resp = r.json() - server_location = resp["data"][0]["location"] - - # Close the file handle - new_file[1].close() - - config.config.set("media", location, f"/files/{server_location}") - - with open(config.config_path, "w+") as f: - config.write(f) + Media.upload(path) def rm(self, path): """Remove local media file from remote server and local config""" diff --git a/ctfcli/core/challenge.py b/ctfcli/core/challenge.py index 3105634..bc33732 100644 --- a/ctfcli/core/challenge.py +++ b/ctfcli/core/challenge.py @@ -19,6 +19,7 @@ RemoteChallengeNotFound, ) from ctfcli.core.image import Image +from ctfcli.core.media import Media from ctfcli.utils.hashing import hash_file from ctfcli.utils.tools import strings @@ -56,6 +57,7 @@ class Challenge(dict): "host", "connection_info", "healthcheck", + "media", "solution", "attempts", "logic", @@ -80,6 +82,7 @@ class Challenge(dict): "topics", "tags", "files", + "media", "hints", "requirements", "state", @@ -294,12 +297,18 @@ def _validate_files(self): if not (self.challenge_directory / challenge_file).exists(): raise InvalidChallengeFile(f"File {challenge_file} could not be loaded") + def _validate_media(self): + media = self.get("media") or [] + for media_file in media: + if not (self.challenge_directory / media_file).exists(): + raise InvalidChallengeFile(f"Media file {media_file} could not be loaded") + def _get_initial_challenge_payload(self, ignore: tuple[str] = ()) -> dict: challenge = self challenge_payload = { "name": self["name"], "category": self.get("category", ""), - "description": self.get("description", ""), + "description": Media.replace_placeholders(self.get("description", "")), "attribution": self.get("attribution", ""), "type": self.get("type", "standard"), # Hide the challenge for the duration of the sync / creation @@ -393,15 +402,14 @@ def _delete_file(self, remote_location: str): r.raise_for_status() def _create_file(self, local_path: Path): - 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={"file": new_file}, data=file_payload) - r.raise_for_status() + with local_path.open(mode="rb") as local_file: + new_file = (local_path.name, local_file) - # Close the file handle - new_file[1].close() + # Specifically use data= here to send multipart/form-data + r = self.api.post("/api/v1/files", files={"file": new_file}, data=file_payload) + r.raise_for_status() def _create_all_files(self): new_files = [] @@ -419,6 +427,11 @@ def _create_all_files(self): for file_payload in new_files: file_payload[1][1].close() + def _upload_media(self): + media = self.get("media") or [] + for media_file in media: + Media.upload(self.challenge_directory / media_file) + def _delete_existing_hints(self): remote_hints = self.api.get("/api/v1/hints").json()["data"] for hint in remote_hints: @@ -589,18 +602,20 @@ def _create_solution(self): for mdx, alt, path in markdown_images: 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={"file": new_file}, data=file_payload) - r.raise_for_status() - resp = r.json() - server_location = resp["data"][0]["location"] - content = content.replace(mdx, f"![{alt}](/files/{server_location})") + with local_path.open(mode="rb") as local_file: + new_file = (local_path.name, local_file) + + # Specifically use data= here to send multipart/form-data + 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"] + content = content.replace(mdx, f"![{alt}](/files/{server_location})") # Process snippet includes (--8<-- "filename") for full_match, filename in snippet_includes: @@ -958,6 +973,11 @@ def sync(self, ignore: tuple[str] = ()) -> None: # _validate_files will raise if file is not found self._validate_files() + if challenge.get("media", False) and "media" not in ignore: + # _validate_media will raise if file is not found + self._validate_media() + self._upload_media() + challenge_payload = self._get_initial_challenge_payload(ignore=ignore) self._load_challenge_id() @@ -1033,9 +1053,8 @@ def sync(self, ignore: tuple[str] = ()) -> None: local_file_sha1sum = hash_file(lf) # Allow users to specify sha1sum in ignore to force reuploads - if "sha1sum" not in ignore: - if local_file_sha1sum == remote_file_sha1sum: - continue + if "sha1sum" not in ignore and local_file_sha1sum == remote_file_sha1sum: + continue # if sha1sums are not present, or the hashes are different, re-upload the file self._delete_file(remote_files[local_file_name]["location"]) @@ -1106,6 +1125,11 @@ def create(self, ignore: tuple[str] = ()) -> None: # _validate_files will raise if file is not found self._validate_files() + if challenge.get("media", False) and "media" not in ignore: + # _validate_media will raise if file is not found + self._validate_media() + self._upload_media() + challenge_payload = self._get_initial_challenge_payload(ignore=ignore) # in the case of creation, value and type can't be ignored: diff --git a/ctfcli/core/media.py b/ctfcli/core/media.py index 59b0de6..b5b4d2a 100644 --- a/ctfcli/core/media.py +++ b/ctfcli/core/media.py @@ -1,14 +1,52 @@ +from pathlib import Path + +from ctfcli.core.api import API from ctfcli.core.config import Config +from ctfcli.core.exceptions import ProjectNotInitialized from ctfcli.utils.tools import safe_format class Media: @staticmethod - def replace_placeholders(content: str) -> str: + def upload(path) -> str: config = Config() + if config.config.has_section("media") is False: + config.config.add_section("media") + + api = API() + + path = Path(path) + filename = path.name + new_file = (filename, path.open(mode="rb")) + location = f"media/{filename}" + file_payload = { + "type": "page", + "location": location, + } + + try: + # Specifically use data= here to send multipart/form-data + 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"] + finally: + new_file[1].close() + + media_url = f"/files/{server_location}" + config.config.set("media", location, media_url) + + with open(config.config_path, "w+") as f: + config.write(f) + + return media_url + + @staticmethod + def replace_placeholders(content: str) -> str: try: + config = Config() section = config["media"] - except KeyError: + except (KeyError, ProjectNotInitialized): section = [] for m in section: content = safe_format(content, items={m: config["media"][m]}) diff --git a/tests/core/test_challenge.py b/tests/core/test_challenge.py index 9dcfcfb..2ee518c 100644 --- a/tests/core/test_challenge.py +++ b/tests/core/test_challenge.py @@ -392,6 +392,19 @@ class TestSyncChallenge(unittest.TestCase): minimal_challenge = BASE_DIR / "fixtures" / "challenges" / "test-challenge-minimal" / "challenge.yml" files_challenge = BASE_DIR / "fixtures" / "challenges" / "test-challenge-files" / "challenge.yml" + @mock.patch("ctfcli.core.challenge.Challenge.load_installed_challenges", return_value=installed_challenges) + @mock.patch("ctfcli.core.challenge.Media.upload") + @mock.patch("ctfcli.core.challenge.API") + def test_uploads_media(self, mock_api_constructor: MagicMock, mock_media_upload: MagicMock, *args, **kwargs): + challenge = Challenge(self.minimal_challenge, {"media": ["challenge.yml"]}) + + mock_api: MagicMock = mock_api_constructor.return_value + challenge.sync() + + mock_media_upload.assert_called_once_with(challenge.challenge_directory / "challenge.yml") + mock_api.post.assert_not_called() + mock_api.delete.assert_not_called() + @mock.patch("ctfcli.core.challenge.Challenge.load_installed_challenges", return_value=installed_challenges) @mock.patch("ctfcli.core.challenge.API") def test_updates_simple_properties(self, mock_api_constructor: MagicMock, *args, **kwargs): @@ -1519,6 +1532,7 @@ def test_does_not_update_ignored_attributes(self): "topics", "tags", "files", + "media", "hints", "requirements", "module", @@ -1598,7 +1612,7 @@ def test_does_not_update_ignored_attributes(self): if p == "extra": challenge["extra"] = {"new-value": "new-value"} - if p in ["flags", "topics", "tags", "files", "hints", "requirements"]: + if p in ["flags", "topics", "tags", "files", "media", "hints", "requirements"]: challenge[p] = ["new-value"] if p == "module": @@ -1625,6 +1639,14 @@ def test_does_not_update_ignored_attributes(self): mock_api.post.assert_not_called() mock_api.delete.assert_not_called() + def test_exits_if_media_do_not_exist(self): + challenge = Challenge(self.minimal_challenge, {"media": ["files/nonexistent.png"]}) + + with self.assertRaises(InvalidChallengeFile) as e: + challenge.sync() + + self.assertEqual(str(e.exception), "Media file files/nonexistent.png could not be loaded") + class TestCreateChallenge(unittest.TestCase): installed_challenges = [ @@ -1657,6 +1679,31 @@ class TestCreateChallenge(unittest.TestCase): minimal_challenge = BASE_DIR / "fixtures" / "challenges" / "test-challenge-minimal" / "challenge.yml" full_challenge = BASE_DIR / "fixtures" / "challenges" / "test-challenge-full" / "challenge.yml" + @mock.patch("ctfcli.core.challenge.Challenge.load_installed_challenges", return_value=installed_challenges) + @mock.patch("ctfcli.core.challenge.Media.upload") + @mock.patch("ctfcli.core.challenge.API") + def test_uploads_media_on_create( + self, mock_api_constructor: MagicMock, mock_media_upload: MagicMock, *args, **kwargs + ): + challenge = Challenge(self.minimal_challenge, {"media": ["challenge.yml"]}) + + def mock_post(*args, **kwargs): + path = args[0] + + if path == "/api/v1/challenges": + mock_response = MagicMock() + mock_response.json.return_value = {"success": True, "data": {"id": 3}} + return mock_response + + return MagicMock() + + mock_api: MagicMock = mock_api_constructor.return_value + mock_api.post.side_effect = mock_post + + challenge.create() + + mock_media_upload.assert_called_once_with(challenge.challenge_directory / "challenge.yml") + @mock.patch("ctfcli.core.challenge.Challenge.load_installed_challenges", return_value=installed_challenges) @mock.patch("ctfcli.core.challenge.API") def test_creates_standard_challenge(self, mock_api_constructor: MagicMock, *args, **kwargs):