Add ability to store manage media files locally and reference their URLs in Pages

Author: ColdHeat

ctfcli/__main__.py

@@ -12,6 +12,7 @@
 from ctfcli.cli.challenges import ChallengeCommand
 from ctfcli.cli.config import ConfigCommand
 from ctfcli.cli.instance import InstanceCommand
+from ctfcli.cli.media import MediaCommand
 from ctfcli.cli.pages import PagesCommand
 from ctfcli.cli.plugins import PluginsCommand
 from ctfcli.cli.templates import TemplatesCommand
@@ -111,6 +112,9 @@ def challenge(self):
     def pages(self):
         return COMMANDS.get("pages")
 
+    def media(self):
+        return COMMANDS.get("media")
+
     def plugins(self):
         return COMMANDS.get("plugins")
 
@@ -125,6 +129,7 @@ def templates(self):
     "plugins": PluginsCommand(),
     "templates": TemplatesCommand(),
     "instance": InstanceCommand(),
+    "media": MediaCommand(),
     "cli": CTFCLI(),
 }
 

ctfcli/cli/media.py

@@ -0,0 +1,55 @@
+import os
+
+from ctfcli.core.api import API
+from ctfcli.core.config import Config
+
+
+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()
+
+        new_file = ("file", open(path, mode="rb"))
+        filename = os.path.basename(path)
+        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.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)
+
+    def rm(self, path):
+        """Remove local media file from remote server and local config"""
+        config = Config()
+        api = API()
+
+        local_location = config["media"][path]
+
+        remote_files = api.get("/api/v1/files?type=page").json()["data"]
+        for remote_file in remote_files:
+            if f"/files/{remote_file['location']}" == local_location:
+                # Delete file from server
+                r = api.delete(f"/api/v1/files/{remote_file['id']}")
+                r.raise_for_status()
+
+                # Update local config file
+                del config["media"][path]
+                with open(config.config_path, "w+") as f:
+                    config.write(f)

ctfcli/core/media.py

@@ -0,0 +1,11 @@
+from ctfcli.core.config import Config
+from ctfcli.utils.tools import safe_format
+
+
+class Media:
+    @staticmethod
+    def replace_placeholders(content: str) -> str:
+        config = Config()
+        for m in config["media"]:
+            content = safe_format(content, items={m: config["media"][m]})
+        return content

ctfcli/core/page.py

@@ -17,6 +17,7 @@
     InvalidPageConfiguration,
     InvalidPageFormat,
 )
+from ctfcli.core.media import Media
 
 PAGE_FORMATS = {
     ".md": "markdown",
@@ -85,8 +86,8 @@ def _get_data_by_path(self) -> Optional[Dict]:
 
         with open(self.page_path, "r") as page_file:
             page_data = frontmatter.load(page_file)
-
-            return {**page_data.metadata, "content": page_data.content}
+            content = Media.replace_placeholders(page_data.content)
+            return {**page_data.metadata, "content": content}
 
     def _get_data_by_id(self) -> Optional[Dict]:
         r = self.api.get(f"/api/v1/pages/{self.page_id}")
@@ -173,6 +174,8 @@ def get_format(ext) -> str:
 
     @staticmethod
     def get_format_extension(fmt) -> str:
+        if fmt is None:
+            return ".md"
         for supported_ext, supported_fmt in PAGE_FORMATS.items():
             if fmt == supported_fmt:
                 return supported_ext

ctfcli/utils/tools.py

@@ -1,3 +1,4 @@
+import re
 import string
 
 
@@ -21,3 +22,11 @@ def strings(filename, min_length=4):
 
         if len(result) >= min_length:  # catch result at EOF
             yield result
+
+
+def safe_format(fmt, items):
+    """
+    Function that safely formats strings with arbitrary potentially user-supplied format strings
+    Looks for interpolation placeholders like {target} or {{ target }}
+    """
+    return re.sub(r"\{?\{([^{}]*)\}\}?", lambda m: items.get(m.group(1).strip(), m.group(0)), fmt)