vishalbhat07/ctfd

1#!/usr/bin/env python2# -*- coding: utf-8 -*-3import zipfile4 5from CTFd.plugins import (6    bypass_csrf_protection,7    get_admin_plugin_menu_bar,8    get_user_page_menu_bar,9    override_function,10    override_template,11    register_admin_plugin_menu_bar,12    register_admin_plugin_script,13    register_admin_plugin_stylesheet,14    register_plugin_asset,15    register_plugin_assets_directory,16    register_plugin_script,17    register_user_page_menu_bar,18)19from CTFd.utils.exports import export_ctf, import_ctf20from tests.helpers import (21    create_ctfd,22    destroy_ctfd,23    gen_challenge,24    login_as_user,25    setup_ctfd,26)27 28 29def test_register_plugin_asset():30    """Test that plugin asset registration works"""31    app = create_ctfd(setup=False)32    register_plugin_asset(app, asset_path="/plugins/__init__.py")33    app = setup_ctfd(app)34    with app.app_context():35        with app.test_client() as client:36            r = client.get("/plugins/__init__.py")37            assert len(r.get_data(as_text=True)) > 038            assert r.status_code == 20039    destroy_ctfd(app)40 41 42def test_register_plugin_assets_directory():43    """Test that plugin asset directory registration works"""44    app = create_ctfd(setup=False)45    register_plugin_assets_directory(app, base_path="/plugins/")46    app = setup_ctfd(app)47    with app.app_context():48        with app.test_client() as client:49            r = client.get("/plugins/__init__.py")50            assert len(r.get_data(as_text=True)) > 051            assert r.status_code == 20052            r = client.get("/plugins/challenges/__init__.py")53            assert len(r.get_data(as_text=True)) > 054            assert r.status_code == 20055    destroy_ctfd(app)56 57 58def test_override_template():59    """Does override_template work properly for regular themes when used from a plugin"""60    app = create_ctfd()61    with app.app_context():62        override_template("login.html", "LOGIN OVERRIDE")63        with app.test_client() as client:64            r = client.get("/login")65            assert r.status_code == 20066            output = r.get_data(as_text=True)67            assert "LOGIN OVERRIDE" in output68    destroy_ctfd(app)69 70 71def test_admin_override_template():72    """Does override_template work properly for the admin panel when used from a plugin"""73    app = create_ctfd()74    with app.app_context():75        override_template("admin/users/user.html", "ADMIN USER OVERRIDE")76 77        client = login_as_user(app, name="admin", password="password")78        r = client.get("/admin/users/1")79        assert r.status_code == 20080        output = r.get_data(as_text=True)81        assert "ADMIN USER OVERRIDE" in output82    destroy_ctfd(app)83 84 85def test_register_plugin_script():86    """Test that register_plugin_script adds script paths to the core theme when used from a plugin"""87    app = create_ctfd()88    with app.app_context():89        register_plugin_script("/fake/script/path.js")90        register_plugin_script("http://examplectf.com/fake/script/path.js")91        with app.test_client() as client:92            r = client.get("/")93            output = r.get_data(as_text=True)94            assert "/fake/script/path.js" in output95            assert "http://examplectf.com/fake/script/path.js" in output96    destroy_ctfd(app)97 98 99def test_register_plugin_stylesheet():100    """Test that register_plugin_stylesheet adds stylesheet paths to the core theme when used from a plugin"""101    app = create_ctfd()102    with app.app_context():103        register_plugin_script("/fake/stylesheet/path.css")104        register_plugin_script("http://examplectf.com/fake/stylesheet/path.css")105        with app.test_client() as client:106            r = client.get("/")107            output = r.get_data(as_text=True)108            assert "/fake/stylesheet/path.css" in output109            assert "http://examplectf.com/fake/stylesheet/path.css" in output110    destroy_ctfd(app)111 112 113def test_register_admin_plugin_script():114    """Test that register_admin_plugin_script adds script paths to the admin theme when used from a plugin"""115    app = create_ctfd()116    with app.app_context():117        register_admin_plugin_script("/fake/script/path.js")118        register_admin_plugin_script("http://examplectf.com/fake/script/path.js")119        with login_as_user(app, name="admin") as client:120            r = client.get("/admin/statistics")121            output = r.get_data(as_text=True)122            assert "/fake/script/path.js" in output123            assert "http://examplectf.com/fake/script/path.js" in output124    destroy_ctfd(app)125 126 127def test_register_admin_plugin_stylesheet():128    """Test that register_admin_plugin_stylesheet adds stylesheet paths to the admin theme when used from a plugin"""129    app = create_ctfd()130    with app.app_context():131        register_admin_plugin_stylesheet("/fake/stylesheet/path.css")132        register_admin_plugin_stylesheet(133            "http://examplectf.com/fake/stylesheet/path.css"134        )135        with login_as_user(app, name="admin") as client:136            r = client.get("/admin/statistics")137            output = r.get_data(as_text=True)138            assert "/fake/stylesheet/path.css" in output139            assert "http://examplectf.com/fake/stylesheet/path.css" in output140    destroy_ctfd(app)141 142 143def test_register_admin_plugin_menu_bar():144    """145    Test that register_admin_plugin_menu_bar() properly inserts into HTML and get_admin_plugin_menu_bar()146    returns the proper list.147    """148    app = create_ctfd()149    with app.app_context():150        register_admin_plugin_menu_bar(151            title="test_admin_plugin_name", route="/test_plugin"152        )153 154        client = login_as_user(app, name="admin", password="password")155        r = client.get("/admin/statistics")156        output = r.get_data(as_text=True)157        assert "/test_plugin" in output158        assert "test_admin_plugin_name" in output159 160        menu_item = get_admin_plugin_menu_bar()[0]161        assert menu_item.title == "test_admin_plugin_name"162        assert menu_item.route == "/test_plugin"163    destroy_ctfd(app)164 165 166def test_register_user_page_menu_bar():167    """168    Test that the register_user_page_menu_bar() properly inserts into HTML and get_user_page_menu_bar() returns the169    proper list.170    """171    app = create_ctfd()172    with app.app_context():173        register_user_page_menu_bar(174            title="test_user_menu_link", route="/test_user_href"175        )176 177        with app.test_client() as client:178            r = client.get("/")179            output = r.get_data(as_text=True)180            assert "/test_user_href" in output181            assert "test_user_menu_link" in output182 183        with app.test_request_context():184            menu_item = get_user_page_menu_bar()[0]185            assert menu_item.title == "test_user_menu_link"186            assert menu_item.route == "/test_user_href"187    destroy_ctfd(app)188 189 190def test_bypass_csrf_protection():191    """192    Test that the bypass_csrf_protection decorator functions properly193    """194    app = create_ctfd()195 196    with app.app_context():197        with app.test_client() as client:198            r = client.post("/login")199            output = r.get_data(as_text=True)200            assert r.status_code == 403201 202        def bypass_csrf_protection_test_route():203            return "Success", 200204 205        # Hijack an existing route to avoid any kind of hacks to create a test route206        app.view_functions["auth.login"] = bypass_csrf_protection(207            bypass_csrf_protection_test_route208        )209 210        with app.test_client() as client:211            r = client.post("/login")212            output = r.get_data(as_text=True)213            assert r.status_code == 200214            assert output == "Success"215    destroy_ctfd(app)216 217 218def test_challenges_model_access_plugin_class():219    """220    Test that the Challenges model can access its plugin class221    """222    app = create_ctfd()223 224    with app.app_context():225        from CTFd.plugins.challenges import get_chal_class226 227        chal = gen_challenge(app.db)228        assert chal.plugin_class == get_chal_class("standard")229    destroy_ctfd(app)230 231 232def test_import_ctf_override():233    """Test that import_ctf can be overridden"""234    app = create_ctfd()235    if not app.config.get("SQLALCHEMY_DATABASE_URI").startswith("sqlite"):236        with app.app_context():237 238            def override_func(backup, *args, **kwargs):239                return "OVERRIDDEN"240 241            override_function("import_ctf", override_func)242            result = import_ctf("dummy_backup")243            assert result == "OVERRIDDEN"244            assert app.overridden_functions["import_ctf"] is override_func245 246            # Test real import_ctf can still be called247            try:248                import_ctf("dummy_backup", ignore_overrides=True)249            except zipfile.BadZipfile:250                # This is expected251                pass252 253    destroy_ctfd(app)254 255 256def test_export_ctf_override():257    """Test that export_ctf can be overridden"""258    app = create_ctfd()259    with app.app_context():260 261        def override_func():262            return "EXPORT_OVERRIDDEN"263 264        override_function("export_ctf", override_func)265        result = export_ctf()266        assert result == "EXPORT_OVERRIDDEN"267        assert app.overridden_functions["export_ctf"] is override_func268 269        # Test real export_ctf can still be called270        result = export_ctf(ignore_overrides=True)271        assert result != "EXPORT_OVERRIDDEN"272        assert hasattr(result, "read")273 274    destroy_ctfd(app)275