__init__.py

""" Plugin entry-point """

import json
import os
import re

from authlib.integrations.flask_client import OAuth
from sqlalchemy import inspect, text

from CTFd.config import process_boolean_str
from CTFd.models import db
from CTFd.plugins import override_template
from CTFd.utils import get_app_config

from .blueprint import load_bp
from .models import OAuthClients

PLUGIN_PATH = os.path.dirname(__file__)
CONFIG = json.load(open("{}/config.json".format(PLUGIN_PATH)))


def migrate_db():
    """Apply incremental schema changes to oauth_clients that create_all() won't add."""
    inspector = inspect(db.engine)
    if 'oauth_clients' not in inspector.get_table_names():
        return  # Table doesn't exist yet; create_all() will build it fresh.
    existing = {col['name'] for col in inspector.get_columns('oauth_clients')}
    migrations = [
        ('server_metadata_url', 'TEXT'),
        ('enabled', 'BOOLEAN NOT NULL DEFAULT TRUE'),
    ]
    with db.engine.begin() as conn:
        for col_name, col_type in migrations:
            if col_name not in existing:
                conn.execute(text(f'ALTER TABLE oauth_clients ADD COLUMN {col_name} {col_type}'))


def oauth_clients():
    return OAuthClients.query.all()


def update_login_template(app):
    """
    Gets the actual login template and injects 
    the SSO buttons before the Forms.auth.LoginForm block
    """

    environment = app.jinja_environment
    original = app.jinja_loader.get_source(environment, 'login.html')[0]

    match = re.search(".*Forms\.auth\.LoginForm.*\n", original)

    # If Forms.auth.LoginForm is not found (maybe in a custom template), it does nothing
    if match:
        pos = match.start()

        PLUGIN_PATH = os.path.dirname(__file__)
        if process_boolean_str(get_app_config("OAUTH_NO_LOCAL_USERS")):
            injecting_file_path = os.path.join(
                PLUGIN_PATH, 'templates/login_oauth_no_local.html')
        else:
             injecting_file_path = os.path.join(
                PLUGIN_PATH, 'templates/login_oauth.html')
        with open(injecting_file_path, 'r') as f:
            injecting = f.read()

        if process_boolean_str(get_app_config("OAUTH_NO_LOCAL_USERS")):
            match2 = re.search("{%\s*endwith\s*%}", original)
            pos2 = match2.end()
            new_template = original[:pos] + injecting + original[pos2:]
        else:
            new_template = original[:pos] + injecting + original[pos:]
        override_template('login.html', new_template)


def update_challenge_template(app):
    """
    Gets the actual challenges template and injects an
    event listener to test is an SSO logout has occurred
    """

    environment = app.jinja_environment
    original = app.jinja_loader.get_source(environment, 'challenges.html')[0]
    match = re.search('{% block scripts %}', original)
    if match:
        pos = match.end()
        injecting_file_path = os.path.join(PLUGIN_PATH, 'templates/challenges_sso.html')
        with open(injecting_file_path, 'r') as f:
            injecting = f.read()
        original = original[:pos] + injecting + original[pos:]

    override_template('challenges.html', original)


def update_settings_template(app):
    """
    Gets the actual settings template and disabled the
    name and email fields as treated by our IDP. If
    OAUTH_NO_LOCAL_USERS is set, disable password as well.

    Yes the user can bypass this, but if they do the impact
    is that they might loose access to CTFd.. Their problem,
    not ours.
    """

    environment = app.jinja_environment
    original = app.jinja_loader.get_source(environment, 'settings.html')[0]
    match = re.search('form.name\(', original)
    if match:
        pos = match.end()
        original = original[:pos] + 'disabled=True, '+ original[pos:]

    match = re.search('form.email\(', original)
    if match:
        pos = match.end()
        original = original[:pos] + 'disabled=True, '+ original[pos:]

    if process_boolean_str(get_app_config("OAUTH_NO_LOCAL_USERS")):
        match = re.search('form.confirm\(', original)
        if match:
            pos = match.end()
            original = original[:pos] + 'disabled=True, '+ original[pos:]

        match = re.search('form.password\(', original)
        if match:
            pos = match.end()
            original = original[:pos] + 'disabled=True, '+ original[pos:]

    override_template('settings.html', original)


def numactive(clients):
  n = 0
  for client in clients:
      if client.enabled:
          n += 1
  return n

def load(app):
    # Create database tables
    app.db.create_all()

    # Apply column-level migrations for existing tables
    migrate_db()

    # Get all saved clients and register them
    clients = oauth_clients()
    oauth = OAuth(app)
    for client in clients:
        client.register(oauth)

    # Register oauth_clients() as template global
    app.jinja_env.globals.update(oauth_clients=oauth_clients)

    # Update the login template
    if process_boolean_str(get_app_config("OAUTH_CREATE_BUTTONS")) or \
       process_boolean_str(get_app_config("OAUTH_NO_LOCAL_USERS")):
        update_login_template(app)
        update_challenge_template(app)
        update_settings_template(app)

    # Register the blueprint containing the routes
    bp = load_bp(oauth)
    app.register_blueprint(bp)

    # Add a function to Jinja2 to count our active Oauth providers
    app.jinja_env.globals.update(numactive=numactive)

    if process_boolean_str(get_app_config("OAUTH_SSO_LOGOUT")):
        # Overwrite existing logout function to treat SSO logout
        app.view_functions["auth.logout"] = app.view_functions["sso.sso_logout"]