From 4e80b514f488ceac7588b7c6c65481af05208578 Mon Sep 17 00:00:00 2001
From: Kevin Chung <kchung@nyu.edu>
Date: Sat, 14 Oct 2017 17:17:06 -0400
Subject: [PATCH] Add plugin wrappers for plugin related utils functions (#410)

---
 .gitignore                 |  3 ++
 CTFd/plugins/__init__.py   | 32 +++++++++++++++++--
 tests/test_plugin_utils.py | 63 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 95 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 11aabc08..0c3c59dd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -66,3 +66,6 @@ CTFd/uploads
 
 # Vagrant
 .vagrant
+
+# CTFd Exports
+*.zip
\ No newline at end of file
diff --git a/CTFd/plugins/__init__.py b/CTFd/plugins/__init__.py
index 0f5c6af4..9457a695 100644
--- a/CTFd/plugins/__init__.py
+++ b/CTFd/plugins/__init__.py
@@ -3,8 +3,13 @@ import importlib
 import os
 
 from flask.helpers import safe_join
-from flask import send_file, send_from_directory, abort
-from CTFd.utils import admins_only as admins_only_wrapper
+from flask import current_app as app, send_file, send_from_directory, abort
+from CTFd.utils import (
+    admins_only as admins_only_wrapper,
+    override_template as utils_override_template,
+    register_plugin_script as utils_register_plugin_script,
+    register_plugin_stylesheet as utils_register_plugin_stylesheet
+)
 
 
 def register_plugin_assets_directory(app, base_path, admins_only=False):
@@ -48,6 +53,29 @@ def register_plugin_asset(app, asset_path, admins_only=False):
     app.add_url_rule(rule=rule, endpoint=asset_path, view_func=asset_handler)
 
 
+def override_template(*args, **kwargs):
+    """
+    Overrides a template with the provided html content.
+
+    e.g. override_template('scoreboard.html', '<h1>scores</h1>')
+    """
+    utils_override_template(*args, **kwargs)
+
+
+def register_plugin_script(*args, **kwargs):
+    """
+    Adds a given script to the base.html template which all pages inherit from
+    """
+    utils_register_plugin_script(*args, **kwargs)
+
+
+def register_plugin_stylesheet(*args, **kwargs):
+    """
+    Adds a given stylesheet to the base.html template which all pages inherit from.
+    """
+    utils_register_plugin_stylesheet(*args, **kwargs)
+
+
 def init_plugins(app):
     """
     Searches for the load function in modules in the CTFd/plugins folder. This function is called with the current CTFd
diff --git a/tests/test_plugin_utils.py b/tests/test_plugin_utils.py
index 63362b8d..942f432a 100644
--- a/tests/test_plugin_utils.py
+++ b/tests/test_plugin_utils.py
@@ -3,7 +3,13 @@
 
 from tests.helpers import *
 from CTFd.models import ip2long, long2ip
-from CTFd.plugins import register_plugin_assets_directory, register_plugin_asset
+from CTFd.plugins import (
+    register_plugin_assets_directory,
+    register_plugin_asset,
+    register_plugin_script,
+    register_plugin_stylesheet,
+    override_template
+)
 from freezegun import freeze_time
 from mock import patch
 import json
@@ -37,3 +43,58 @@ def test_register_plugin_assets_directory():
             assert len(r.get_data(as_text=True)) > 0
             assert r.status_code == 200
     destroy_ctfd(app)
+
+
+def test_override_template():
+    """Does override_template work properly for regular themes when used from a plugin"""
+    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
+    destroy_ctfd(app)
+
+
+def test_admin_override_template():
+    """Does override_template work properly for the admin panel when used from a plugin"""
+    app = create_ctfd()
+    with app.app_context():
+        override_template('admin/team.html', 'ADMIN TEAM OVERRIDE')
+
+        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
+    destroy_ctfd(app)
+
+
+def test_register_plugin_script():
+    '''Test that register_plugin_script adds script paths to the original theme when used from a plugin'''
+    app = create_ctfd()
+    with app.app_context():
+        register_plugin_script('/fake/script/path.js')
+        register_plugin_script('http://ctfd.io/fake/script/path.js')
+        with app.test_client() as client:
+            r = client.get('/')
+            output = r.get_data(as_text=True)
+            assert '/fake/script/path.js' in output
+            assert 'http://ctfd.io/fake/script/path.js' in output
+    destroy_ctfd(app)
+
+
+def test_register_plugin_stylesheet():
+    '''Test that register_plugin_stylesheet adds stylesheet paths to the original theme when used from a plugin'''
+    app = create_ctfd()
+    with app.app_context():
+        register_plugin_script('/fake/stylesheet/path.css')
+        register_plugin_script('http://ctfd.io/fake/stylesheet/path.css')
+        with app.test_client() as client:
+            r = client.get('/')
+            output = r.get_data(as_text=True)
+            assert '/fake/stylesheet/path.css' in output
+            assert 'http://ctfd.io/fake/stylesheet/path.css' in output
+    destroy_ctfd(app)