Repository: CTFd/CTFd
Pull request #297 (2017-06-28 01:05:31)

Adding override_template function
Purely for use by plugins to replace templates in an existing theme


Patch:
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


Test patch:
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