diff --git a/CTFd/constants/__init__.py b/CTFd/constants/__init__.py new file mode 100644 index 00000000..8f58d2d7 --- /dev/null +++ b/CTFd/constants/__init__.py @@ -0,0 +1,63 @@ +from enum import Enum +from flask import current_app + +JS_ENUMS = {} + + +class RawEnum(Enum): + """ + This is a customized enum class which should be used with a mixin. + The mixin should define the types of each member. + + For example: + + class Colors(str, RawEnum): + RED = "red" + GREEN = "green" + BLUE = "blue" + """ + + def __str__(self): + return str(self._value_) + + @classmethod + def keys(cls): + return list(cls.__members__.keys()) + + @classmethod + def values(cls): + return list(cls.__members__.values()) + + @classmethod + def test(cls, value): + try: + return bool(cls(value)) + except ValueError: + return False + + +def JSEnum(cls): + """ + This is a decorator used to gather all Enums which should be shared with + the CTFd front end. The JS_Enums dictionary can be taken be a script and + compiled into a JavaScript file for use by frontend assets. JS_Enums + should not be passed directly into Jinja. A JinjaEnum is better for that. + """ + if cls.__name__ not in JS_ENUMS: + JS_ENUMS[cls.__name__] = dict(cls.__members__) + else: + raise KeyError("{} was already defined as a JSEnum".format(cls.__name__)) + return cls + + +def JinjaEnum(cls): + """ + This is a decorator used to inject the decorated Enum into Jinja globals + which allows you to access it from the front end. If you need to access + an Enum from JS, a better tool to use is the JSEnum decorator. + """ + if cls.__name__ not in current_app.jinja_env.globals: + current_app.jinja_env.globals[cls.__name__] = cls + else: + raise KeyError("{} was already defined as a JinjaEnum".format(cls.__name__)) + return cls diff --git a/manage.py b/manage.py index b388b94b..0e8e8cd5 100644 --- a/manage.py +++ b/manage.py @@ -12,6 +12,21 @@ manager = Manager(app) manager.add_command("db", MigrateCommand) +def jsenums(): + from CTFd.constants import JS_ENUMS + import json + import os + + path = os.path.join(app.root_path, "themes/core/assets/js/constants.js") + + with open(path, "w+") as f: + for k, v in JS_ENUMS.items(): + f.write("const {} = Object.freeze({});".format(k, json.dumps(v))) + + +BUILD_COMMANDS = {"jsenums": jsenums} + + @manager.command def get_config(key): with app.app_context(): @@ -24,5 +39,12 @@ def set_config(key, value): print(set_config_util(key, value).value) +@manager.command +def build(cmd): + with app.app_context(): + cmd = BUILD_COMMANDS.get(cmd) + cmd() + + if __name__ == "__main__": manager.run() diff --git a/tests/constants/test_constants.py b/tests/constants/test_constants.py new file mode 100644 index 00000000..c16d2294 --- /dev/null +++ b/tests/constants/test_constants.py @@ -0,0 +1,52 @@ +from CTFd.constants import RawEnum, JSEnum, JinjaEnum +from tests.helpers import create_ctfd, destroy_ctfd + + +def test_RawEnum(): + class Colors(str, RawEnum): + RED = "red" + GREEN = "green" + BLUE = "blue" + + class Numbers(str, RawEnum): + ONE = 1 + TWO = 2 + THREE = 3 + + assert Colors.RED == "red" + assert Colors.GREEN == "green" + assert Colors.BLUE == "blue" + assert Colors.test("red") is True + assert Colors.test("purple") is False + assert str(Numbers.ONE) == "1" + assert sorted(Colors.keys()) == sorted(["RED", "GREEN", "BLUE"]) + assert sorted(Colors.values()) == sorted(["red", "green", "blue"]) + + +def test_JSEnum(): + from CTFd.constants import JS_ENUMS + import json + + @JSEnum + class Colors(str, RawEnum): + RED = "red" + GREEN = "green" + BLUE = "blue" + + assert JS_ENUMS["Colors"] == {"RED": "red", "GREEN": "green", "BLUE": "blue"} + assert json.dumps(JS_ENUMS) + + +def test_JinjaEnum(): + app = create_ctfd() + with app.app_context(): + + @JinjaEnum + class Colors(str, RawEnum): + RED = "red" + GREEN = "green" + BLUE = "blue" + + assert app.jinja_env.globals["Colors"] is Colors + assert app.jinja_env.globals["Colors"].RED == "red" + destroy_ctfd(app)