Override image

dojo_plugin/models/__init__.py

@@ -361,7 +361,7 @@ class DojoChallenges(db.Model):
     description = db.Column(db.Text)
 
     data = db.Column(db.JSON)
-    data_fields = ["path_override"]
+    data_fields = ["image_override", "path_override"]
 
     dojo = db.relationship("Dojos",
                            foreign_keys=[dojo_id],
@@ -387,7 +387,7 @@ def __init__(self, *args, **kwargs):
             if kwargs.get("challenge") is not None:
                 raise AttributeError("Import requires challenge to be None")
 
-            for field in ["id", "name", "description", "challenge"]:
+            for field in ["id", "name", "description", "challenge", "image_override"]:
                 kwargs[field] = kwargs[field] if kwargs.get(field) is not None else getattr(default, field, None)
 
             # TODO: maybe we should track the entire import
@@ -472,7 +472,7 @@ def path(self):
 
     @property
     def image(self):
-        return "pwncollege-challenge"
+        return self.image_override or "pwncollege-challenge"
 
     def challenge_paths(self, user):
         secret = current_app.config["SECRET_KEY"]

dojo_plugin/utils/dojo/__init__.py

@@ -13,7 +13,7 @@
 from schema import Schema, Optional, Regex, Or, Use, SchemaError
 from flask import abort, g
 from sqlalchemy.orm.exc import NoResultFound
-from CTFd.models import db, Users, Challenges, Flags, Solves
+from CTFd.models import db, Users, Challenges, Flags, Solves, Admins
 from CTFd.utils.user import get_current_user, is_admin
 
 from ...models import Dojos, DojoUsers, DojoModules, DojoChallenges, DojoResources, DojoChallengeVisibilities, DojoResourceVisibilities
@@ -23,6 +23,7 @@
 ID_REGEX = Regex(r"^[a-z0-9-]{1,32}$")
 UNIQUE_ID_REGEX = Regex(r"^[a-z0-9-~]{1,128}$")
 NAME_REGEX = Regex(r"^[\S ]{1,128}$")
+IMAGE_REGEX = Regex(r"^[\S]{1,256}$")
 DATE = Use(datetime.datetime.fromisoformat)
 
 ID_NAME_DESCRIPTION = {
@@ -49,6 +50,8 @@
         Optional("emoji"): Regex(r"^\S$"),
     },
 
+    Optional("image"): IMAGE_REGEX,
+
     Optional("import"): {
         "dojo": UNIQUE_ID_REGEX,
     },
@@ -57,6 +60,8 @@
         **ID_NAME_DESCRIPTION,
         **VISIBILITY,
 
+        Optional("image"): IMAGE_REGEX,
+
         Optional("import"): {
             Optional("dojo"): UNIQUE_ID_REGEX,
             "module": ID_REGEX,
@@ -66,7 +71,7 @@
             **ID_NAME_DESCRIPTION,
             **VISIBILITY,
 
-            # Optional("image", default="pwncollege-challenge"): Regex(r"^[\S ]{1, 256}$"),
+            Optional("image"): IMAGE_REGEX,
             # Optional("path"): Regex(r"^[^\s\.\/][^\s\.]{,255}$"),
 
             Optional("import"): {
@@ -215,24 +220,26 @@ def visibility(cls, *args):
         if start or stop:
             return cls(start=start, stop=stop)
 
-    def import_ids(ids, *datas):
-        results = {
-            id: None
-            for id in ids
-        }
-        for data in datas:
-            for id, result in results.items():
-                results[id] = data.get("import", {}).get(id, None) or result
-        for id, result in results.items():
-            assert result is not None, f"Missing `{id}` in import"
-        return tuple(results.values())
+    _missing = object()
+    def shadow(attr, *datas, default=_missing):
+        for data in reversed(datas):
+            if attr in data:
+                return data[attr]
+        if default is not _missing:
+            return default
+        raise KeyError(f"Missing `{attr}` in `{datas}`")
+
+    def import_ids(attrs, *datas):
+        datas_import = [data.get("import", {}) for data in datas]
+        return tuple(shadow(id, *datas_import) for id in attrs)
 
     dojo.modules = [
         DojoModules(
             **{kwarg: module_data.get(kwarg) for kwarg in ["id", "name", "description"]},
             challenges=[
                 DojoChallenges(
                     **{kwarg: challenge_data.get(kwarg) for kwarg in ["id", "name", "description"]},
+                    image=shadow("image", dojo_data, module_data, challenge_data, default=None),
                     challenge=challenge(module_data.get("id"), challenge_data.get("id")) if "import" not in challenge_data else None,
                     visibility=visibility(DojoChallengeVisibilities, dojo_data, module_data, challenge_data),
                     default=(assert_one(DojoChallenges.from_id(*import_ids(["dojo", "module", "challenge"], dojo_data, module_data, challenge_data)),
@@ -286,6 +293,10 @@ def import_ids(ids, *datas):
                 students = yaml.safe_load(students_yml_path.read_text())
                 dojo.course["students"] = students
 
+    custom_image = any(challenge.image for challenge in dojo.challenges)
+    admin_dojo = any(isinstance(dojo_admin.user, Admins) for dojo_admin in dojo.admins)
+    assert not (custom_image and not admin_dojo), "Custom images are only allowed for admin dojos"
+
     return dojo