add media for challenge descriptions

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

ctfcli/core/challenge.py

@@ -20,6 +20,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
 
@@ -57,6 +58,7 @@ class Challenge(dict):
         "host",
         "connection_info",
         "healthcheck",
+        "media",
         "solution",
         "attempts",
         "logic",
@@ -82,6 +84,7 @@ class Challenge(dict):
         "topics",
         "tags",
         "files",
+        "media",
         "hints",
         "requirements",
         "scheduled_at",
@@ -365,12 +368,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
@@ -469,15 +478,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 = []
@@ -495,6 +503,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:
@@ -665,18 +678,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:
@@ -1037,6 +1052,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()
@@ -1112,9 +1132,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"])
@@ -1185,6 +1204,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:

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]})

tests/core/test_challenge.py

@@ -393,6 +393,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):
@@ -1535,6 +1548,7 @@ def test_does_not_update_ignored_attributes(self):
             "topics",
             "tags",
             "files",
+            "media",
             "hints",
             "requirements",
             "module",
@@ -1620,7 +1634,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":
@@ -1647,6 +1661,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 = [
@@ -1679,6 +1701,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):