diff --git a/CTFd/__init__.py b/CTFd/__init__.py index d1b32c9b..e60cc35f 100644 --- a/CTFd/__init__.py +++ b/CTFd/__init__.py @@ -21,11 +21,22 @@ __version__ = '1.0.2' class ThemeLoader(FileSystemLoader): + def __init__(self, searchpath, encoding='utf-8', followlinks=False): + super(ThemeLoader, self).__init__(searchpath, encoding, followlinks) + self.overriden_templates = {} + def get_source(self, environment, template): + # Check if the template has been overriden + if template in self.overriden_templates: + return self.overriden_templates[template], template, True + + # Check if the template requested is for the admin panel if template.startswith('admin/'): template = template.lstrip('admin/') template = "/".join(['admin', 'templates', template]) return super(ThemeLoader, self).get_source(environment, template) + + # Load regular theme data theme = utils.get_config('ctf_theme') template = "/".join([theme, 'templates', template]) return super(ThemeLoader, self).get_source(environment, template) diff --git a/CTFd/utils.py b/CTFd/utils.py index 2d6f21c9..f29866f8 100644 --- a/CTFd/utils.py +++ b/CTFd/utils.py @@ -166,6 +166,11 @@ def hide_scores(): return get_config('hide_scores') +def override_template(template, html): + with app.app_context(): + app.jinja_loader.overriden_templates[template] = html + + def pages(): pages = Pages.query.filter(Pages.route != "index").all() return pages diff --git a/tests/test_utils.py b/tests/test_utils.py index dd2e5407..2ea8dad5 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,6 +3,7 @@ from tests.helpers import * from CTFd.models import ip2long, long2ip +from CTFd.utils import override_template import json @@ -26,3 +27,28 @@ def test_long2ip_ipv6(): """Does long2ip work properly for ipv6 addresses""" assert long2ip(42540766452641154071740215577757643572) == '2001:db8:85a3::8a2e:370:7334' assert long2ip(42540616829182469433547762482097946625) == '2001:658:22a:cafe:200::1' + + +def test_override_template(): + """Does override_template work properly for regular themes""" + app = create_ctfd() + with app.app_context(): + override_template('login.html', 'LOGIN OVERRIDE') + with app.test_client() as client: + r = client.get('/login') + assert r.status_code == 200 + output = r.get_data(as_text=True) + assert 'LOGIN OVERRIDE' in output + + +def test_admin_override_template(): + """Does override_template work properly for the admin panel""" + app = create_ctfd() + with app.app_context(): + override_template('admin/team.html', 'ADMIN TEAM OVERRIDE') + with app.app_context(): + client = login_as_user(app, name="admin", password="password") + r = client.get('/admin/team/1') + assert r.status_code == 200 + output = r.get_data(as_text=True) + assert 'ADMIN TEAM OVERRIDE' in output