23880703624

Committed 02 Apr 2026 02:26AM UTC coverage: 90.624% (+0.01%) from 90.614%

Build # 23880703624

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #1195 from scoringengine/feature/inject-categories-1174

Add inject category field (Business/Technical/Incident Response)

Coverage Stats

7 of 9 new or added lines in 3 files covered. (77.78%)

3 existing lines in 1 file now uncovered.

4736 of 5226 relevant lines covered (90.62%)

0.91 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

87.71

/scoring_engine/web/views/api/admin.py

import html
import json
import os
import re
import time
from datetime import datetime, timezone

import pytz
from dateutil.parser import parse
from flask import flash, jsonify, redirect, request, url_for
from flask_login import current_user, login_required


def _ensure_utc_aware(dt):
    """Ensure datetime is timezone-aware in UTC. Handles both naive and aware datetimes."""
    if dt is None:
        return None
    if dt.tzinfo is None:
        # Naive datetime - assume UTC
        return pytz.utc.localize(dt)
    # Already aware - convert to UTC
    return dt.astimezone(pytz.utc)


from sqlalchemy.orm import joinedload
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.sql import func

from scoring_engine.cache_helper import (
    update_all_cache,
    update_all_inject_data,
    update_inject_comments,
    update_inject_data,
    update_overview_data,
    update_scoreboard_data,
    update_service_data,
    update_services_data,
    update_services_navbar,
    update_team_stats,
)
from scoring_engine.events import publish_event
from scoring_engine.celery_stats import CeleryStats
from scoring_engine.config import config
from scoring_engine.db import db
from scoring_engine.engine.basic_check import CHECK_FAILURE_TEXT, CHECK_SUCCESS_TEXT, CHECK_TIMED_OUT_TEXT
from scoring_engine.engine.engine import Engine
from scoring_engine.engine.execute_command import execute_command
from scoring_engine.models.check import Check
from scoring_engine.models.environment import Environment
from scoring_engine.models.inject import Inject, InjectComment, InjectRubricScore, RubricItem, Template
from scoring_engine.models.kb import KB
from scoring_engine.models.property import Property
from scoring_engine.models.round import Round
from scoring_engine.models.service import Service
from scoring_engine.models.setting import Setting
from scoring_engine.models.team import Team
from scoring_engine.models.user import User
from scoring_engine.models.welcome import get_welcome_config, save_welcome_config
from scoring_engine.notifications import notify_inject_graded, notify_revision_requested

from . import mod


@mod.route("/api/admin/update_environment_info", methods=["POST"])
@login_required
def admin_update_environment():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            environment = db.session.get(Environment, int(request.form["pk"]))
            if environment:
                if request.form["name"] in ("matching_content", "matching_content_reject"):
                    value = html.escape(request.form["value"])
                    if value:
                        try:
                            re.compile(value)
                        except re.error as e:
                            return jsonify({"error": "Invalid regex pattern: " + str(e)}), 400
                    setattr(environment, request.form["name"], value or None)
                db.session.add(environment)
                db.session.commit()
                return jsonify({"status": "Updated Environment Information"})
            return jsonify({"error": "Incorrect permissions"})
    return jsonify({"error": "Incorrect permissions"})


@mod.route("/api/admin/update_property", methods=["POST"])
@login_required
def admin_update_property():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            property_obj = db.session.get(Property, int(request.form["pk"]))
            if property_obj:
                if request.form["name"] == "property_name":
                    property_obj.name = html.escape(request.form["value"])
                elif request.form["name"] == "property_value":
                    property_obj.value = html.escape(request.form["value"])
                db.session.add(property_obj)
                db.session.commit()
                return jsonify({"status": "Updated Property Information"})
            return jsonify({"error": "Incorrect permissions"})
    return jsonify({"error": "Incorrect permissions"})


@mod.route("/api/admin/update_check", methods=["POST"])
@login_required
def admin_update_check():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            check = db.session.get(Check, int(request.form["pk"]))
            if check:
                modified_check = False
                if request.form["name"] == "check_value":
                    if request.form["value"] == "1":
                        check.result = True
                    elif request.form["value"] == "2":
                        check.result = False
                    modified_check = True
                elif request.form["name"] == "check_reason":
                    modified_check = True
                    check.reason = request.form["value"]
                if modified_check:
                    db.session.add(check)
                    db.session.commit()
                    update_scoreboard_data()
                    update_overview_data()
                    update_services_navbar(check.service.team.id)
                    update_team_stats(check.service.team.id)
                    update_services_data(check.service.team.id)
                    update_service_data(check.service.id)
                    return jsonify({"status": "Updated Property Information"})
            return jsonify({"error": "Incorrect permissions"})
    return jsonify({"error": "Incorrect permissions"})


@mod.route("/api/admin/update_host", methods=["POST"])
@login_required
def admin_update_host():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            service = db.session.get(Service, int(request.form["pk"]))
            if service:
                if request.form["name"] == "host":
                    service.host = html.escape(request.form["value"])
                    db.session.add(service)
                    db.session.commit()
                    update_overview_data()
                    update_services_data(service.team.id)
                    update_service_data(service.id)
                    return jsonify({"status": "Updated Service Information"})
    return jsonify({"error": "Incorrect permissions"})


@mod.route("/api/admin/update_port", methods=["POST"])
@login_required
def admin_update_port():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            service = db.session.get(Service, int(request.form["pk"]))
            if service:
                if request.form["name"] == "port":
                    service.port = int(request.form["value"])
                    db.session.add(service)
                    db.session.commit()
                    update_overview_data()
                    update_services_data(service.team.id)
                    update_service_data(service.id)
                    return jsonify({"status": "Updated Service Information"})
    return jsonify({"error": "Incorrect permissions"})


@mod.route("/api/admin/update_worker_queue", methods=["POST"])
@login_required
def admin_update_worker_queue():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            service = db.session.get(Service, int(request.form["pk"]))
            if service:
                if request.form["name"] == "worker_queue":
                    service.worker_queue = request.form["value"]
                    db.session.add(service)
                    db.session.commit()
                    return jsonify({"status": "Updated Service Information"})
    return jsonify({"error": "Incorrect permissions"})


@mod.route("/api/admin/update_points", methods=["POST"])
@login_required
def admin_update_points():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            service = db.session.get(Service, int(request.form["pk"]))
            if service:
                if request.form["name"] == "points":
                    service.points = int(request.form["value"])
                    db.session.add(service)
                    db.session.commit()
                    return jsonify({"status": "Updated Service Information"})
    return jsonify({"error": "Incorrect permissions"})


@mod.route("/api/admin/update_about_page_content", methods=["POST"])
@login_required
def admin_update_about_page_content():
    if current_user.is_white_team:
        if "about_page_content" in request.form:
            setting = Setting.get_setting("about_page_content")
            setting.value = request.form["about_page_content"]
            db.session.add(setting)
            db.session.commit()
            Setting.clear_cache("about_page_content")
            flash("About Page Content Successfully Updated.", "success")
            return redirect(url_for("admin.settings"))
        flash("Error: about_page_content not specified.", "danger")
        return redirect(url_for("admin.manage"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_welcome_page_content", methods=["POST"])
@login_required
def admin_update_welcome_page_content():
    if current_user.is_white_team:
        if "welcome_page_content" in request.form:
            setting = Setting.get_setting("welcome_page_content")
            setting.value = request.form["welcome_page_content"]
            db.session.add(setting)
            db.session.commit()
            Setting.clear_cache("welcome_page_content")
            flash("Welcome Page Content Successfully Updated.", "success")
            return redirect(url_for("admin.settings"))
        flash("Error: welcome_page_content not specified.", "danger")
        return redirect(url_for("admin.manage"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_target_round_time", methods=["POST"])
@login_required
def admin_update_target_round_time():
    if current_user.is_white_team:
        if "target_round_time" in request.form:
            setting = Setting.get_setting("target_round_time")
            input_time = request.form["target_round_time"]
            if not input_time.isdigit():
                flash("Error: Target Round Time must be an integer.", "danger")
                return redirect(url_for("admin.settings"))
            setting.value = input_time
            db.session.add(setting)
            db.session.commit()
            Setting.clear_cache("target_round_time")
            flash("Target Round Time Successfully Updated.", "success")
            return redirect(url_for("admin.settings"))
        flash("Error: target_round_time not specified.", "danger")
        return redirect(url_for("admin.settings"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_worker_refresh_time", methods=["POST"])
@login_required
def admin_update_worker_refresh_time():
    if current_user.is_white_team:
        if "worker_refresh_time" in request.form:
            setting = Setting.get_setting("worker_refresh_time")
            input_time = request.form["worker_refresh_time"]
            if not input_time.isdigit():
                flash("Error: Worker Refresh Time must be an integer.", "danger")
                return redirect(url_for("admin.settings"))
            setting.value = input_time
            db.session.add(setting)
            db.session.commit()
            Setting.clear_cache("worker_refresh_time")
            flash("Worker Refresh Time Successfully Updated.", "success")
            return redirect(url_for("admin.settings"))
        flash("Error: worker_refresh_time not specified.", "danger")
        return redirect(url_for("admin.settings"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_blueteam_edit_hostname", methods=["POST"])
@login_required
def admin_update_blueteam_edit_hostname():
    if current_user.is_white_team:
        setting = Setting.get_setting("blue_team_update_hostname")
        if setting.value is True:
            setting.value = False
        else:
            setting.value = True
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("blue_team_update_hostname")
        return redirect(url_for("admin.permissions"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_blueteam_edit_port", methods=["POST"])
@login_required
def admin_update_blueteam_edit_port():
    if current_user.is_white_team:
        setting = Setting.get_setting("blue_team_update_port")
        if setting.value is True:
            setting.value = False
        else:
            setting.value = True
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("blue_team_update_port")
        return redirect(url_for("admin.permissions"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_blueteam_edit_account_usernames", methods=["POST"])
@login_required
def admin_update_blueteam_edit_account_usernames():
    if current_user.is_white_team:
        setting = Setting.get_setting("blue_team_update_account_usernames")
        if setting.value is True:
            setting.value = False
        else:
            setting.value = True
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("blue_team_update_account_usernames")
        return redirect(url_for("admin.permissions"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_blueteam_edit_account_passwords", methods=["POST"])
@login_required
def admin_update_blueteam_edit_account_passwords():
    if current_user.is_white_team:
        setting = Setting.get_setting("blue_team_update_account_passwords")
        if setting.value is True:
            setting.value = False
        else:
            setting.value = True
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("blue_team_update_account_passwords")
        return redirect(url_for("admin.permissions"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_blueteam_view_check_output", methods=["POST"])
@login_required
def admin_update_blueteam_view_check_output():
    if current_user.is_white_team:
        setting = Setting.get_setting("blue_team_view_check_output")
        if setting.value is True:
            setting.value = False
        else:
            setting.value = True
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("blue_team_view_check_output")
        return redirect(url_for("admin.permissions"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_anonymize_team_names", methods=["POST"])
@login_required
def admin_update_anonymize_team_names():
    if current_user.is_white_team:
        setting = Setting.get_setting("anonymize_team_names")
        if setting is None:
            # Create the setting if it doesn't exist
            setting = Setting(name="anonymize_team_names", value=True)
        elif setting.value is True:
            setting.value = False
        else:
            setting.value = True
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("anonymize_team_names")
        return redirect(url_for("admin.permissions"))
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/check/<int:check_id>/full_output")
@login_required
def admin_get_check_full_output(check_id):
    if current_user.is_white_team:
        check = db.session.get(Check, check_id)
        if not check:
            return jsonify({"error": "Check not found"}), 404

        team_name = check.service.team.name
        service_name = check.service.name
        round_num = check.round.number

        # Try disk file first, fall back to DB output
        output_path = os.path.join(
            config.check_output_folder,
            team_name,
            service_name,
            f"round_{round_num}.txt",
        )

        if os.path.isfile(output_path):
            with open(output_path, "r") as f:
                content = f.read()
        else:
            content = check.output or ""

        return content, 200, {"Content-Type": "text/plain; charset=utf-8"}
    return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/get_round_progress")
@login_required
def get_check_progress_total():
    if current_user.is_white_team:
        task_id_settings = db.session.query(KB).filter_by(name="task_ids").order_by(KB.round_num.desc()).first()
        total_stats = {}
        total_stats["finished"] = 0
        total_stats["pending"] = 0

        team_stats = {}
        if task_id_settings:
            task_dict = json.loads(task_id_settings.value)
            for team_name, task_ids in task_dict.items():
                for task_id in task_ids:
                    task = execute_command.AsyncResult(task_id)
                    if team_name not in team_stats:
                        team_stats[team_name] = {}
                        team_stats[team_name]["pending"] = 0
                        team_stats[team_name]["finished"] = 0

                    if task.state == "PENDING":
                        team_stats[team_name]["pending"] += 1
                        total_stats["pending"] += 1
                    else:
                        team_stats[team_name]["finished"] += 1
                        total_stats["finished"] += 1

        total_percentage = 0
        total_tasks = total_stats["finished"] + total_stats["pending"]
        if total_stats["finished"] == 0:
            total_percentage = 0
        elif total_tasks == 0:
            total_percentage = 100
        elif total_stats and total_stats["finished"]:
            total_percentage = int((total_stats["finished"] / total_tasks) * 100)

        output_dict = {"Total": total_percentage}
        for team_name, team_stat in team_stats.items():
            team_total_percentage = 0
            team_total_tasks = team_stat["finished"] + team_stat["pending"]
            if team_stat["finished"] == 0:
                team_total_percentage = 0
            elif team_total_tasks == 0:
                team_total_percentage = 100
            elif team_stat and team_stat["finished"]:
                team_total_percentage = int((team_stat["finished"] / team_total_tasks) * 100)
            output_dict[team_name] = team_total_percentage

        return json.dumps(output_dict)
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/inject/<inject_id>/reopen", methods=["POST"])
@login_required
def admin_post_inject_reopen(inject_id):
    if current_user.is_white_team:
        inject = db.session.get(Inject, inject_id)
        if not inject:
            return jsonify({"status": "Invalid Inject ID"}), 400
        if inject.status != "Graded":
            return jsonify({"status": "Inject is not in Graded state"}), 400

        inject.status = "Submitted"
        inject.graded = None
        comment = InjectComment(
            f"Inject reopened for re-grading by {current_user.username}.",
            current_user,
            inject,
        )
        db.session.add(comment)
        db.session.commit()
        update_inject_data(inject_id)
        update_inject_comments(inject_id)
        update_scoreboard_data()
        publish_event("inject_update", {"inject_id": inject_id}, visibility="blue", team_id=inject.team_id)
        publish_event("inject_update", {"inject_id": inject_id}, visibility="white")
        return jsonify({"status": "Success"}), 200
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/injects/templates/<template_id>")
@login_required
def admin_get_inject_templates_id(template_id):
    if current_user.is_white_team:
        template = db.session.get(Template, int(template_id))
        data = dict(
            id=template.id,
            title=template.title,
            scenario=template.scenario,
            deliverable=template.deliverable,
            category=template.category,
            max_score=template.max_score,
            start_time=_ensure_utc_aware(template.start_time).astimezone(pytz.timezone(config.timezone)).isoformat(),
            end_time=_ensure_utc_aware(template.end_time).astimezone(pytz.timezone(config.timezone)).isoformat(),
            enabled=template.enabled,
            rubric_items=[
                {"id": x.id, "title": x.title, "description": x.description, "points": x.points, "order": x.order}
                for x in template.rubric_items
            ],
            teams=[inject.team.name for inject in template.injects if inject.enabled and inject.team],
        )
        return jsonify(data)
    else:
        return jsonify({"status": "Unauthorized"}), 403


@mod.route("/api/admin/injects/templates/<template_id>", methods=["PUT"])
@login_required
def admin_put_inject_templates_id(template_id):
    if current_user.is_white_team:
        data = request.get_json()
        template = db.session.get(Template, int(template_id))
        if template:
            if data.get("title"):
                template.title = data["title"]
            if data.get("scenario"):
                template.scenario = data["scenario"]
            if data.get("deliverable"):
                template.deliverable = data["deliverable"]
            if data.get("start_time"):
                template.start_time = parse(data["start_time"]).astimezone(pytz.utc).replace(tzinfo=None)
            if data.get("end_time"):
                template.end_time = parse(data["end_time"]).astimezone(pytz.utc).replace(tzinfo=None)
            if "category" in data:
                template.category = data["category"] or None
            # TODO - Fix this to not be string values from javascript select
            if data.get("status") == "Enabled":
                template.enabled = True
            elif data.get("status") == "Disabled":
                template.enabled = False
            if "rubric_items" in data:
                # Update existing rubric items in place to preserve scores
                existing_by_id = {item.id: item for item in template.rubric_items}
                incoming_ids = set()
                for idx, ri in enumerate(data["rubric_items"]):
                    ri_id = ri.get("id")
                    if ri_id and ri_id in existing_by_id:
                        # Update existing item
                        item = existing_by_id[ri_id]
                        item.title = ri["title"]
                        item.points = ri["points"]
                        item.description = ri.get("description")
                        item.order = ri.get("order", idx)
                        incoming_ids.add(ri_id)
                    else:
                        # Create new item
                        rubric_item = RubricItem(
                            title=ri["title"],
                            points=ri["points"],
                            template=template,
                            description=ri.get("description"),
                            order=ri.get("order", idx),
                        )
                        db.session.add(rubric_item)
                # Delete only items that were removed from the payload
                for item_id, item in existing_by_id.items():
                    if item_id not in incoming_ids:
                        db.session.delete(item)
            if data.get("selectedTeams"):
                for team_name in data["selectedTeams"]:
                    inject = (
                        db.session.query(Inject)
                        .join(Template)
                        .join(Team)
                        .filter(Team.name == team_name)
                        .filter(Template.id == template_id)
                        .one_or_none()
                    )
                    # Update inject if it exists
                    if inject:
                        inject.enabled = True
                    # Otherwise, create the inject
                    else:
                        team = db.session.query(Team).filter(Team.name == team_name).first()
                        inject = Inject(
                            team=team,
                            template=template,
                        )
                        db.session.add(inject)
            if data.get("unselectedTeams"):
                injects = (
                    db.session.query(Inject)
                    .join(Template)
                    .join(Team)
                    .filter(Team.name.in_(data["unselectedTeams"]))
                    .filter(Template.id == template_id)
                    .all()
                )
                for inject in injects:
                    inject.enabled = False
            db.session.commit()
            return jsonify({"status": "Success"}), 200
        else:
            return jsonify({"status": "Error", "message": "Template not found"}), 400

    else:
        return jsonify({"status": "Unauthorized"}), 403


@mod.route("/api/admin/injects/templates/<template_id>", methods=["DELETE"])
@login_required
def admin_delete_inject_templates_id(template_id):
    if current_user.is_white_team:
        template = db.session.get(Template, int(template_id))
        if template:
            db.session.delete(template)
            db.session.commit()
            return jsonify({"status": "Success"}), 200
        else:
            return jsonify({"status": "Error", "message": "Template not found"}), 400
    else:
        return jsonify({"status": "Unauthorized"}), 403


@mod.route("/api/admin/injects/templates")
@login_required
def admin_get_inject_templates():
    if current_user.is_white_team:
        data = list()
        templates = db.session.query(Template).options(joinedload(Template.injects)).all()
        for template in templates:
            data.append(
                dict(
                    id=template.id,
                    title=template.title,
                    category=template.category,
                    scenario=template.scenario,
                    deliverable=template.deliverable,
                    max_score=template.max_score,
                    start_time=template.start_time.astimezone(pytz.timezone(config.timezone)).isoformat(),
                    end_time=template.end_time.astimezone(pytz.timezone(config.timezone)).isoformat(),
                    enabled=template.enabled,
                    rubric_items=[
                        {
                            "id": x.id,
                            "title": x.title,
                            "description": x.description,
                            "points": x.points,
                            "order": x.order,
                        }
                        for x in template.rubric_items
                    ],
                    teams=[inject.team.name for inject in template.injects if inject if inject.enabled if inject.team],
                )
            )
        return jsonify(data=data)
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/inject/<inject_id>/grade", methods=["POST"])
@login_required
def admin_post_inject_grade(inject_id):
    if current_user.is_white_team:
        data = request.get_json()
        if "rubric_scores" not in data or not data["rubric_scores"]:
            return jsonify({"status": "Invalid Score Provided"}), 400
        inject = db.session.get(Inject, inject_id)
        if not inject:
            return jsonify({"status": "Invalid Inject ID"}), 400

        # Validate rubric items belong to the template
        template_item_ids = {item.id for item in inject.template.rubric_items}
        template_item_max = {item.id: item.points for item in inject.template.rubric_items}
        for rs in data["rubric_scores"]:
            if rs["rubric_item_id"] not in template_item_ids:
                return jsonify({"status": "Invalid rubric item ID"}), 400
            if rs["score"] > template_item_max[rs["rubric_item_id"]]:
                return jsonify({"status": "Score exceeds rubric item max"}), 400

        # Delete existing scores and create new ones
        for old_score in list(inject.rubric_scores):
            db.session.delete(old_score)
        db.session.flush()

        for rs in data["rubric_scores"]:
            rubric_item = db.session.get(RubricItem, rs["rubric_item_id"])
            score_obj = InjectRubricScore(
                score=rs["score"],
                inject=inject,
                rubric_item=rubric_item,
                grader=current_user,
            )
            db.session.add(score_obj)

        inject.graded = datetime.now(timezone.utc)
        inject.status = "Graded"
        score = sum(rs["score"] for rs in data["rubric_scores"])
        comment = InjectComment(
            f"Inject graded by {current_user.username}. Score: {score}/{inject.template.max_score}.",
            current_user,
            inject,
        )
        db.session.add(comment)
        db.session.commit()
        update_inject_data(inject_id)
        update_inject_comments(inject_id)
        update_scoreboard_data()
        publish_event("inject_update", {"inject_id": inject_id}, visibility="blue", team_id=inject.team_id)
        publish_event("inject_update", {"inject_id": inject_id}, visibility="white")
        notify_inject_graded(inject)
        return jsonify({"status": "Success"}), 200
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/inject/<inject_id>/request-revision", methods=["POST"])
@login_required
def admin_post_inject_request_revision(inject_id):
    if current_user.is_white_team:
        inject = db.session.get(Inject, inject_id)
        if not inject:
            return jsonify({"status": "Invalid Inject ID"}), 400
        if inject.status not in ("Submitted", "Resubmitted"):
            return jsonify({"status": "Inject is not in a submittable state"}), 400

        data = request.get_json() or {}
        reason = data.get("reason", "Revision requested by grader.")

        inject.status = "Revision Requested"
        # Create a system comment with the reason
        comment = InjectComment(content=reason, user=current_user, inject=inject)
        db.session.add(comment)
        db.session.commit()
        update_inject_data(inject_id)
        update_inject_comments(inject_id)
        publish_event("inject_update", {"inject_id": inject_id}, visibility="blue", team_id=inject.team_id)
        publish_event("inject_update", {"inject_id": inject_id}, visibility="white")
        notify_revision_requested(inject)
        return jsonify({"status": "Success"}), 200
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/injects/template/<int:template_id>/submissions")
@login_required
def admin_get_template_submissions(template_id):
    if current_user.is_white_team:
        template = db.session.get(Template, template_id)
        if not template:
            return jsonify({"status": "Template not found"}), 404

        submissions = []
        for inject in template.injects:
            if not inject.team or not inject.enabled:
                continue
            submissions.append(
                {
                    "id": inject.id,
                    "team": inject.team.name,
                    "status": inject.status,
                    "score": inject.score,
                    "max_score": template.max_score,
                    "submitted": inject.submitted.isoformat() if inject.submitted else None,
                    "graded": inject.graded.isoformat() if inject.graded else None,
                }
            )
        return jsonify(data=submissions), 200
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/injects/templates", methods=["POST"])
@login_required
def admin_post_inject_templates():
    if current_user.is_white_team:
        data = request.get_json()
        if (
            data.get("title")
            and data.get("scenario")
            and data.get("deliverable")
            and data.get("start_time")
            and data.get("end_time")
        ):
            template = Template(
                title=data["title"],
                scenario=data["scenario"],
                deliverable=data["deliverable"],
                start_time=parse(data["start_time"]).astimezone(pytz.utc).replace(tzinfo=None),
                end_time=parse(data["end_time"]).astimezone(pytz.utc).replace(tzinfo=None),
                category=data.get("category"),
            )
            db.session.add(template)
            db.session.flush()
            # Create rubric items
            if data.get("rubric_items"):
                for i, item_data in enumerate(data["rubric_items"]):
                    rubric_item = RubricItem(
                        title=item_data["title"],
                        description=item_data.get("description", ""),
                        points=item_data["points"],
                        order=item_data.get("order", i),
                        template=template,
                    )
                    db.session.add(rubric_item)
            db.session.commit()
            # TODO - Fix this to not be string values from javascript select
            if data.get("status") == "Enabled":
                template.enabled = True
            elif data.get("status") == "Disabled":
                template.enabled = False
            if data.get("selectedTeams"):
                for team_name in data["selectedTeams"]:
                    inject = (
                        db.session.query(Inject)
                        .join(Template)
                        .join(Team)
                        .filter(Team.name == team_name)
                        .filter(Template.id == template.id)
                        .one_or_none()
                    )
                    # Update inject if it exists
                    if inject:
                        inject.enabled = True
                    # Otherwise, create the inject
                    else:
                        team = db.session.query(Team).filter(Team.name == team_name).first()
                        inject = Inject(
                            team=team,
                            template=template,
                        )
                        db.session.add(inject)
            if data.get("unselectedTeams"):
                injects = (
                    db.session.query(Inject)
                    .join(Template)
                    .join(Team)
                    .filter(Team.name.in_(data["unselectedTeams"]))
                    .filter(Template.id == template.id)
                    .all()
                )
                for inject in injects:
                    inject.enabled = False
            db.session.commit()
            return jsonify({"status": "Success"}), 200
        else:
            return jsonify({"status": "Error", "message": "Missing Data"}), 400
    else:
        return {"status": "Unauthorized"}, 403


# @mod.route("/api/admin/injects/templates/export")
# @login_required
# def admin_export_inject_templates():
#     if current_user.is_white_team:
#         data = []
#         templates = db.session.query(Template).all()
#         for template in templates:
#             data.append(
#                 dict(
#                     id=template.id,
#                     title=template.title,
#                     scenario=template.scenario,
#                     deliverable=template.deliverable,
#                     start_time=template.start_time,
#                     end_time=template.end_time,
#                     enabled=template.enabled,
#                     rubric=[
#                         {"id": x.id, "value": x.value, "deliverable": x.deliverable}
#                         for x in template.rubric
#                     ],
#                     # TODO - export teams
#                 )
#             )
#         return jsonify(data=data)
#     else:
#         return {"status": "Unauthorized"}, 403


# TODO - Generate injects from templates
@mod.route("/api/admin/injects/templates/import", methods=["POST"])
@login_required
def admin_import_inject_templates():
    if current_user.is_white_team:
        data = request.get_json()
        if data:
            for d in data:
                if d.get("id"):
                    template_id = d["id"]
                    t = db.session.get(Template, int(template_id))
                    # Update template if it exists
                    if t:
                        if d.get("title"):
                            t.title = d["title"]
                        if d.get("scenario"):
                            t.scenario = d["scenario"]
                        if d.get("deliverable"):
                            t.deliverable = d["deliverable"]
                        if d.get("start_time"):
                            t.start_time = parse(d["start_time"]).astimezone(pytz.utc).replace(tzinfo=None)
                        if d.get("end_time"):
                            t.end_time = parse(d["end_time"]).astimezone(pytz.utc).replace(tzinfo=None)
                        if "category" in d:
                            t.category = d["category"] or None
                        if d.get("enabled"):
                            t.enabled = True
                        else:
                            t.enabled = False
                        # for rubric in d["rubric"]:
                        #     if rubric.get("id"):
                        #         rubric_id = rubric["id"]
                        #     r = db.session.get(Rubric, int(rubric_id))
                        #     # Update rubric if it exists
                        #     if r:
                        #         if rubric.get("value"):
                        #             r.value = rubric["value"]
                        #         if rubric.get("deliverable"):
                        #             r.deliverable = rubric["deliverable"]
                        #     # Otherwise, create the rubric
                        #     else:
                        #         r = Rubric(
                        #             value=rubric["value"],
                        #             deliverable=rubric["deliverable"],
                        #             template=template,
                        #         )
                        #         db.session.add(r)
                        # Generate injects from template
                        if d.get("selectedTeams"):
                            for team_name in data["selectedTeams"]:
                                inject = (
                                    db.session.query(Inject)
                                    .join(Template)
                                    .join(Team)
                                    .filter(Team.name == team_name)
                                    .filter(Template.id == template_id)
                                    .one_or_none()
                                )
                                # Update inject if it exists
                                if inject:
                                    inject.enabled = True
                                # Otherwise, create the inject
                                else:
                                    team = db.session.query(Team).filter(Team.name == team_name).first()
                                    inject = Inject(
                                        team=team,
                                        template=template,
                                    )
                                    db.session.add(inject)
                        if d.get("unselectedTeams"):
                            injects = (
                                db.session.query(Inject)
                                .join(Template)
                                .join(Team)
                                .filter(Team.name.in_(data["unselectedTeams"]))
                                .filter(Template.id == template_id)
                                .all()
                            )
                            for inject in injects:
                                inject.enabled = False

                    else:
                        # Template ID doesn't exist, fall through to create a new one
                        d.pop("id")
                if not d.get("id"):
                    # Create the template
                    t = Template(
                        title=d["title"],
                        scenario=d["scenario"],
                        deliverable=d["deliverable"],
                        start_time=parse(d["start_time"]).astimezone(pytz.utc).replace(tzinfo=None),
                        end_time=parse(d["end_time"]).astimezone(pytz.utc).replace(tzinfo=None),
                        enabled=d["enabled"],
                    )
                    db.session.add(t)
                    db.session.flush()
                    # Create rubric items (backward compat: if only "score" exists, create default item)
                    if d.get("rubric_items"):
                        for idx, item_data in enumerate(d["rubric_items"]):
                            rubric_item = RubricItem(
                                title=item_data["title"],
                                description=item_data.get("description", ""),
                                points=item_data["points"],
                                order=item_data.get("order", idx),
                                template=t,
                            )
                            db.session.add(rubric_item)
                    elif d.get("score"):
                        rubric_item = RubricItem(
                            title="Overall Quality",
                            description="",
                            points=int(d["score"]),
                            order=0,
                            template=t,
                        )
                        db.session.add(rubric_item)
                    for team_name in d["teams"]:
                        inject = (
                            db.session.query(Inject)
                            .join(Template)
                            .join(Team)
                            .filter(Team.name == team_name)
                            .filter(Template.id == t.id)
                            .one_or_none()
                        )
                        # Update inject if it exists
                        if inject:
                            inject.enabled = True
                        # Otherwise, create the inject
                        else:
                            team = db.session.query(Team).filter(Team.name == team_name).first()
                            if team:
                                inject = Inject(
                                    team=team,
                                    template=t,
                                )
                                db.session.add(inject)
            db.session.commit()
            return jsonify({"status": "Success"}), 200
        else:
            return jsonify({"status": "Error", "message": "Invalid Data"}), 400
    else:
        return {"status": "Unauthorized"}, 403


# TODO - This is way too many database queries
@mod.route("/api/admin/injects/scores")
@login_required
def admin_inject_scores():
    if current_user.is_white_team:
        data = {}

        injects = (
            db.session.query(Inject)
            .options(joinedload(Inject.template), joinedload(Inject.team))
            .order_by(Inject.template_id)
            .order_by(Inject.team_id)
            .all()
        )

        for inject in injects:
            if inject.team is None:
                continue
            if inject.template.id not in data:
                data[inject.template.id] = {
                    "title": inject.template.title,
                    "end_time": inject.template.end_time,
                }
            if inject.team.name not in data[inject.template.id]:
                data[inject.template.id][inject.team.name] = {
                    "id": inject.id,
                    "score": inject.score,
                    "status": inject.status,
                    "max_score": inject.template.max_score,
                }

        # Rewrite data to be in the format expected by the frontend
        score_data = []
        for x in data.items():
            score_data.append(x[1])

        return jsonify(data=score_data), 200
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/injects/get_bar_chart")
@login_required
def admin_injects_bar():
    if current_user.is_white_team:
        inject_scores = dict(
            db.session.query(Inject.team_id, func.sum(InjectRubricScore.score))
            .join(InjectRubricScore)
            .filter(Inject.status == "Graded")
            .group_by(Inject.team_id)
            .all()
        )

        team_data = {}
        team_labels = []
        team_inject_scores = []
        blue_teams = db.session.query(Team).filter(Team.color == "Blue").order_by(Team.name).all()
        for blue_team in blue_teams:
            team_labels.append(blue_team.name)
            team_inject_scores.append(str(inject_scores.get(blue_team.id, 0)))

        team_data["labels"] = team_labels
        team_data["inject_scores"] = team_inject_scores

        return jsonify(team_data), 200
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/admin_update_template", methods=["POST"])
@login_required
def admin_update_template():
    if current_user.is_white_team:
        if "name" in request.form and "value" in request.form and "pk" in request.form:
            template = db.session.get(Template, int(request.form["pk"]))
            if template:
                modified_check = False
                if request.form["name"] == "template_state":
                    template.state = request.form["value"]
                    modified_check = True
                elif request.form["name"] == "template_points":
                    template.points = request.form["value"]
                    modified_check = True
                if modified_check:
                    db.session.add(template)
                    db.session.commit()
                    # update_scoreboard_data()
                    # update_overview_data()
                    # update_services_navbar(check.service.team.id)
                    # update_team_stats(check.service.team.id)
                    # update_services_data(check.service.team.id)
                    # update_service_data(check.service.id)
                    return jsonify({"status": "Updated Property Information"})
            return jsonify({"error": "Template Not Found"})
    return jsonify({"error": "Incorrect permissions"})


# @mod.route("/api/admin/injects/team/<team_id>")
# @login_required
# def admin_get_team_injects(team_id):
#     if current_user.is_white_team:
#         injects = db.session.query(Inject).filter(team_id == team_id).all()
#         return jsonify(data=injects)
#     else:
#         return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/get_teams")
@login_required
def admin_get_teams():
    if current_user.is_white_team:
        all_teams = db.session.query(Team).all()
        data = []
        for team in all_teams:
            users = {}
            for user in team.users:
                users[user.username] = [user.password, str(user.authenticated).title()]
            data.append({"name": team.name, "color": team.color, "users": users})
        return jsonify(data=data)
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/update_password", methods=["POST"])
@login_required
def admin_update_password():
    if current_user.is_white_team:
        if "user_id" in request.form and "password" in request.form:
            try:
                user_obj = db.session.query(User).filter(User.id == request.form["user_id"]).one()
            except NoResultFound:
                return redirect(url_for("auth.login"))
            user_obj.update_password(html.escape(request.form["password"]))
            user_obj.authenticated = False
            db.session.add(user_obj)
            db.session.commit()
            flash("Password Successfully Updated.", "success")
            return redirect(url_for("admin.manage"))
        else:
            flash("Error: user_id or password not specified.", "danger")
            return redirect(url_for("admin.manage"))
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/add_user", methods=["POST"])
@login_required
def admin_add_user():
    if current_user.is_white_team:
        if "username" in request.form and "password" in request.form and "team_id" in request.form:
            team_obj = db.session.query(Team).filter(Team.id == request.form["team_id"]).one()
            user_obj = User(
                username=html.escape(request.form["username"]),
                password=html.escape(request.form["password"]),
                team=team_obj,
            )
            db.session.add(user_obj)
            db.session.commit()
            flash("User successfully added.", "success")
            return redirect(url_for("admin.manage"))
        else:
            flash("Error: Username, Password, or Team ID not specified.", "danger")
            return redirect(url_for("admin.manage"))
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/add_team", methods=["POST"])
@login_required
def admin_add_team():
    if current_user.is_white_team:
        if "name" in request.form and "color" in request.form:
            team_obj = Team(html.escape(request.form["name"]), html.escape(request.form["color"]))
            db.session.add(team_obj)
            db.session.commit()
            flash("Team successfully added.", "success")
            return redirect(url_for("admin.manage"))
        else:
            flash("Error: Team name or color not defined.", "danger")
            return redirect(url_for("admin.manage"))
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/toggle_inject_scores_visible", methods=["POST"])
@login_required
def admin_toggle_inject_scores_visible():
    if current_user.is_white_team:
        Setting.clear_cache("inject_scores_visible")
        setting = Setting.get_setting("inject_scores_visible")
        setting.value = not setting.value
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("inject_scores_visible")
        update_scoreboard_data()
        update_all_inject_data()
        update_overview_data()
        publish_event("settings_changed", {"setting": "inject_scores_visible"})
        return {"status": "Success"}
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/toggle_engine", methods=["POST"])
@login_required
def admin_toggle_engine():
    if current_user.is_white_team:
        Setting.clear_cache("engine_paused")
        setting = Setting.get_setting("engine_paused")
        setting.value = not setting.value
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("engine_paused")
        publish_event("settings_changed", {"setting": "engine_paused"})
        return {"status": "Success"}
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/get_engine_stats")
@login_required
def admin_get_engine_stats():
    if current_user.is_white_team:
        engine_stats = {}
        engine_stats["round_number"] = Round.get_last_round_num()
        engine_stats["num_passed_checks"] = db.session.query(Check).filter_by(result=True).count()
        engine_stats["num_failed_checks"] = db.session.query(Check).filter_by(result=False).count()
        engine_stats["total_checks"] = db.session.query(Check).count()
        return jsonify(engine_stats)
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/get_engine_paused")
@login_required
def admin_get_engine_status():
    if current_user.is_white_team:
        return jsonify({"paused": Setting.get_setting("engine_paused").value})
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/get_worker_stats")
@login_required
def admin_get_worker_stats():
    if current_user.is_white_team:
        worker_stats = CeleryStats.get_worker_stats()
        return jsonify(data=worker_stats)
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/get_queue_stats")
@login_required
def admin_get_queue_stats():
    if current_user.is_white_team:
        queue_stats = CeleryStats.get_queue_stats()
        return jsonify(data=queue_stats)
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/get_competition_summary")
@login_required
def admin_get_competition_summary():
    if current_user.is_white_team:
        blue_teams = db.session.query(Team).filter(Team.color == "Blue").all()
        total_services = db.session.query(Service).join(Team).filter(Team.color == "Blue").count()
        total_checks = db.session.query(Check).count()
        passed_checks = db.session.query(Check).filter_by(result=True).count()

        overall_uptime = 0.0
        if total_checks > 0:
            overall_uptime = round((passed_checks / total_checks) * 100, 1)

        last_round = Round.get_last_round_num()
        currently_passing = 0
        if last_round > 0:
            currently_passing = (
                db.session.query(Check).join(Round).filter(Round.number == last_round, Check.result == True).count()
            )

        return jsonify(
            {
                "blue_teams": len(blue_teams),
                "total_services": total_services,
                "currently_passing": currently_passing,
                "overall_uptime": overall_uptime,
            }
        )
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/welcome/config", methods=["GET"])
@login_required
def admin_get_welcome_config():
    """Get the welcome page configuration."""
    if current_user.is_white_team:
        config = get_welcome_config()
        return jsonify(data=config)
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/welcome/config", methods=["PUT"])
@login_required
def admin_update_welcome_config():
    """Update the welcome page configuration."""
    if current_user.is_white_team:
        data = request.get_json()
        if not data:
            return (
                jsonify(
                    {
                        "status": "Error",
                        "message": "No data provided",
                    }
                ),
                400,
            )

        try:
            config = save_welcome_config(data)
            return jsonify({"status": "Success", "data": config})
        except ValueError as e:
            return jsonify({"status": "Error", "message": str(e)}), 400
    else:
        return {"status": "Unauthorized"}, 403


@mod.route("/api/admin/welcome/upload_logo", methods=["POST"])
@login_required
def admin_upload_sponsor_logo():
    """Upload a sponsor logo image."""
    import os

    from werkzeug.utils import secure_filename

    if not current_user.is_white_team:
        return jsonify({"status": "Unauthorized"}), 403

    if "file" not in request.files:
        return (
            jsonify(
                {
                    "status": "Error",
                    "message": "No file provided",
                }
            ),
            400,
        )

    file = request.files["file"]
    if file.filename == "":
        return (
            jsonify(
                {
                    "status": "Error",
                    "message": "No file selected",
                }
            ),
            400,
        )

    # Validate file extension
    allowed = {"png", "jpg", "jpeg", "gif", "svg", "webp"}
    ext = file.filename.rsplit(".", 1)[-1].lower() if "." in file.filename else ""
    if ext not in allowed:
        return (
            jsonify(
                {
                    "status": "Error",
                    "message": "File type not allowed. Use: " + ", ".join(sorted(allowed)),
                }
            ),
            400,
        )

    filename = secure_filename(file.filename)

    # Save to the shared upload folder (Docker volume), not the static dir.
    # Static files are served by nginx from a separate container/mount,
    # so uploads to the web container's static/ dir would be inaccessible.
    sponsors_dir = os.path.join(config.upload_folder, "sponsors")
    os.makedirs(sponsors_dir, exist_ok=True)

    # Avoid overwriting: add suffix if file exists
    filepath = os.path.join(sponsors_dir, filename)
    if os.path.exists(filepath):
        name, dot_ext = os.path.splitext(filename)
        counter = 1
        while os.path.exists(filepath):
            filename = f"{name}_{counter}{dot_ext}"
            filepath = os.path.join(sponsors_dir, filename)
            counter += 1

    file.save(filepath)

    url = f"/api/admin/welcome/sponsor_logo/{filename}"
    return jsonify({"status": "Success", "url": url})


@mod.route("/api/admin/welcome/sponsor_logo/<filename>")
def serve_sponsor_logo(filename):
    """Serve a sponsor logo from the upload folder."""
    import os

    from flask import abort, send_from_directory
    from werkzeug.utils import secure_filename

    safe_filename = secure_filename(filename)
    if not safe_filename:
        abort(404)
    sponsors_dir = os.path.join(config.upload_folder, "sponsors")
    return send_from_directory(sponsors_dir, safe_filename)


# Global cache for loaded check classes (loaded once per app lifecycle)
_check_classes = None


def _get_check_classes():
    """Load and cache all check classes."""
    global _check_classes
    if _check_classes is None:
        _check_classes = {}
        loaded_checks = Engine.load_check_files(config.checks_location)
        for check_class in loaded_checks:
            _check_classes[check_class.__name__] = check_class
    return _check_classes


@mod.route("/api/admin/check/dry_run", methods=["POST"])
@login_required
def admin_check_dry_run():
    """
    Run a service check without recording results to the database.
    Dispatches to a Celery worker so the check runs in an environment
    that has the required tools (SSH client, DNS utils, etc.).
    """
    if not current_user.is_white_team:
        return jsonify({"status": "Unauthorized"}), 403

    data = request.get_json()
    if not data or "service_id" not in data:
        return jsonify({"status": "error", "message": "service_id is required"}), 400

    service_id = data.get("service_id")
    environment_id = data.get("environment_id")

    service = db.session.get(Service, int(service_id))
    if not service:
        return jsonify({"status": "error", "message": f"Service with id {service_id} not found"}), 404

    # Get environment (specific or random)
    if environment_id:
        environment = db.session.get(Environment, int(environment_id))
        if not environment or environment.service_id != service.id:
            return jsonify({"status": "error", "message": f"Environment {environment_id} not found for service"}), 404
    else:
        if not service.environments:
            return jsonify({"status": "error", "message": "Service has no environments configured"}), 400
        import random

        environment = random.choice(service.environments)

    # Load check class and generate command
    check_classes = _get_check_classes()
    check_class = check_classes.get(service.check_name)
    if not check_class:
        return jsonify({"status": "error", "message": f"Check class '{service.check_name}' not found"}), 404

    try:
        check_obj = check_class(environment)
        command = check_obj.command()
    except LookupError as e:
        return jsonify({"status": "error", "message": f"Check configuration error: {str(e)}"}), 400
    except Exception as e:
        return jsonify({"status": "error", "message": f"Error instantiating check: {str(e)}"}), 500

    # Dispatch to a Celery worker via the service's worker queue
    import time

    job = {"command": command}
    start_time = time.time()
    try:
        result = execute_command.apply_async(args=[job], queue=service.worker_queue)
        completed_job = result.get(timeout=35)
    except Exception as e:
        return jsonify({"status": "error", "message": f"Worker execution failed: {str(e)}"}), 500
    execution_time_ms = int((time.time() - start_time) * 1000)

    # Evaluate result against matching_content
    if completed_job.get("errored_out"):
        check_result = False
        reason = CHECK_TIMED_OUT_TEXT
    else:
        output = completed_job.get("output", "")
        if re.search(environment.matching_content, output):
            check_result = True
            reason = CHECK_SUCCESS_TEXT
        else:
            check_result = False
            reason = CHECK_FAILURE_TEXT

    return jsonify(
        {
            "status": "success",
            "result": check_result,
            "reason": reason,
            "output": completed_job.get("output", "")[:35000],
            "command": command,
            "execution_time_ms": execution_time_ms,
            "service_name": service.name,
            "team_name": service.team.name,
            "host": service.host,
            "port": service.port,
            "matching_content": environment.matching_content,
            "environment_id": environment.id,
        }
    )


@mod.route("/api/admin/rollback", methods=["POST"])
@login_required
def admin_rollback():
    """
    Rollback competition data to a specific round.

    Deletes all rounds, checks, and KB entries from the specified round onwards.
    This is a destructive operation - use with caution.

    Request JSON:
    {
        "round_number": 10,  // Delete this round and all after it
        "confirm": true      // Required to confirm destructive action
    }
    """
    if not current_user.is_white_team:
        return jsonify({"status": "Unauthorized"}), 403

    data = request.get_json()
    if not data:
        return jsonify({"status": "error", "message": "Request body required"}), 400

    round_number = data.get("round_number")
    confirm = data.get("confirm", False)

    if round_number is None:
        return jsonify({"status": "error", "message": "round_number is required"}), 400

    try:
        round_number = int(round_number)
    except (ValueError, TypeError):
        return jsonify({"status": "error", "message": "round_number must be an integer"}), 400

    if round_number < 1:
        return jsonify({"status": "error", "message": "round_number must be >= 1"}), 400

    if not confirm:
        return jsonify({"status": "error", "message": "confirm=true required for destructive operation"}), 400

    # Get current round count for validation
    current_round = Round.get_last_round_num()
    if current_round == 0:
        return jsonify({"status": "error", "message": "No rounds exist to rollback"}), 400

    if round_number > current_round:
        return jsonify(
            {"status": "error", "message": f"round_number ({round_number}) exceeds current round ({current_round})"}
        ), 400

    # Pause the engine to prevent race conditions during rollback
    was_paused = Setting.get_setting("engine_paused").value
    if not was_paused:
        setting = Setting.get_setting("engine_paused")
        setting.value = True
        db.session.add(setting)
        db.session.commit()
        Setting.clear_cache("engine_paused")
        # Wait for the engine to finish its current round and notice the pause
        time.sleep(5)

    try:
        # Revoke any pending Celery tasks for rounds being rolled back
        task_kb_entries = (
            db.session.query(KB)
            .filter(KB.round_num >= round_number, KB.name == "task_ids")
            .all()
        )
        revoked_count = 0
        for kb_entry in task_kb_entries:
            try:
                task_dict = json.loads(kb_entry.value)
                for team_name, task_id_list in task_dict.items():
                    for task_id in task_id_list:
                        execute_command.AsyncResult(task_id).revoke(terminate=True)
                        revoked_count += 1
            except (json.JSONDecodeError, TypeError):
                pass

        # Get rounds to delete
        rounds_to_delete = db.session.query(Round).filter(Round.number >= round_number).all()
        round_ids = [r.id for r in rounds_to_delete]

        # Count checks to delete
        checks_count = db.session.query(Check).filter(Check.round_id.in_(round_ids)).count()

        # Count KB entries to delete
        kb_count = db.session.query(KB).filter(KB.round_num >= round_number).count()

        # Delete checks in batches to avoid lock wait timeout
        BATCH_SIZE = 500
        for i in range(0, len(round_ids), BATCH_SIZE):
            batch_ids = round_ids[i : i + BATCH_SIZE]
            db.session.query(Check).filter(Check.round_id.in_(batch_ids)).delete(synchronize_session=False)
            db.session.commit()

        # Delete KB entries
        db.session.query(KB).filter(KB.round_num >= round_number).delete(synchronize_session=False)
        db.session.commit()

        # Delete rounds
        rounds_count = len(rounds_to_delete)
        db.session.query(Round).filter(Round.number >= round_number).delete(synchronize_session=False)
        db.session.commit()

    finally:
        # Unpause engine if we paused it (even on error)
        if not was_paused:
            Setting.clear_cache("engine_paused")
            setting = Setting.get_setting("engine_paused")
            setting.value = False
            db.session.add(setting)
            db.session.commit()
            Setting.clear_cache("engine_paused")

    # Clear all caches to force recalculation
    from flask import current_app

    update_all_cache(current_app)

    return jsonify(
        {
            "status": "success",
            "message": f"Rolled back to before round {round_number}",
            "deleted": {
                "rounds": rounds_count,
                "checks": checks_count,
                "kb_entries": kb_count,
            },
            "new_current_round": Round.get_last_round_num(),
        }
    )


@mod.route("/api/admin/rollback/preview", methods=["POST"])
@login_required
def admin_rollback_preview():
    """Preview what would be deleted by a rollback operation."""
    if not current_user.is_white_team:
        return jsonify({"status": "Unauthorized"}), 403

    data = request.get_json()
    if not data:
        return jsonify({"status": "error", "message": "Request body required"}), 400

    round_number = data.get("round_number")
    if round_number is None:
        return jsonify({"status": "error", "message": "round_number is required"}), 400

    try:
        round_number = int(round_number)
    except (ValueError, TypeError):
        return jsonify({"status": "error", "message": "round_number must be an integer"}), 400

    if round_number < 1:
        return jsonify({"status": "error", "message": "round_number must be >= 1"}), 400

    current_round = Round.get_last_round_num()
    if current_round == 0:
        return jsonify(
            {
                "status": "success",
                "current_round": 0,
                "round_number": round_number,
                "will_delete": {"rounds": 0, "checks": 0, "kb_entries": 0},
            }
        )

    # Get rounds that would be deleted
    rounds_to_delete = db.session.query(Round).filter(Round.number >= round_number).all()
    round_ids = [r.id for r in rounds_to_delete]

    # Count checks
    checks_count = db.session.query(Check).filter(Check.round_id.in_(round_ids)).count() if round_ids else 0

    # Count KB entries
    kb_count = db.session.query(KB).filter(KB.round_num >= round_number).count()

    return jsonify(
        {
            "status": "success",
            "current_round": current_round,
            "round_number": round_number,
            "will_delete": {
                "rounds": len(rounds_to_delete),
                "checks": checks_count,
                "kb_entries": kb_count,
            },
        }
    )

scoringengine / scoringengine / 23880703624

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous