server.py

import os
import docker
import tarfile
import tempfile
import zipfile
import threading
import csv
import json
from flask import Flask, request, jsonify, send_from_directory
from werkzeug.utils import secure_filename

app = Flask(__name__, static_folder="site")
client = docker.from_env()

# Configuration
UPLOAD_FOLDER = "submissions"
ALLOWED_EXTENSIONS = {"zip", "tar.gz"}
SECRET_TOKEN = "YOUR_SECRET_TOKEN"  # Change this to a secure random string
SITE_DIR = "site"
COOKIE_FILE = os.path.join(UPLOAD_FOLDER, "cookies.json")

app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
csv_lock = threading.Lock()


def get_cookies():
    if not os.path.exists(COOKIE_FILE):
        return {}
    with open(COOKIE_FILE) as f:
        return json.load(f)


def check_cookie(student_id, cookie):
    if not cookie:
        return False
    cookies = get_cookies()
    return cookies.get(student_id) == cookie


def allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS


def run_and_log_result(submission_path, homework_name, student_id):
    """
    This function runs in a background thread.
    It executes the container, waits for the result, and logs it to a CSV file.
    """
    try:
        volumes = {
            os.path.abspath(submission_path): {"bind": "/submission", "mode": "rw"},
            os.path.abspath("/tmp/run.sh"): {
                "bind": "/submission/run.sh",
                "mode": "ro",
            },
        }

        container = client.containers.run(
            "homework-runner",
            command=["./run.sh"],
            volumes=volumes,
            working_dir="/submission",
            detach=True,
        )

        result = container.wait()
        exit_code = result["StatusCode"]
        container.remove()

        # Log the result to the CSV file
        homework_dir = os.path.join(app.config["UPLOAD_FOLDER"], homework_name)
        os.makedirs(homework_dir, exist_ok=True)
        csv_path = os.path.join(homework_dir, "result.csv")

        with csv_lock:
            file_exists = os.path.isfile(csv_path)
            with open(csv_path, "a", newline="") as csvfile:
                writer = csv.writer(csvfile)
                if not file_exists:
                    writer.writerow(["student_id", "exit_code"])
                writer.writerow([student_id, exit_code])

    except Exception as e:
        print(f"Error in background thread for {student_id}/{homework_name}: {e}")


@app.route("/submit/<homework_name>", methods=["POST"])
def submit_homework(homework_name):
    if "file" not in request.files:
        return jsonify({"error": "No file part"}), 400

    file = request.files["file"]
    student_id = request.form.get("student_id")
    cookie = request.form.get("password")

    if not student_id:
        return jsonify({"error": "No student_id provided"}), 400

    if not check_cookie(student_id, cookie):
        return jsonify({"error": "Invalid cookie"}), 401

    if file.filename == "":
        return jsonify({"error": "No selected file"}), 400

    if file and allowed_file(file.filename):
        student_dir = os.path.join(
            app.config["UPLOAD_FOLDER"], homework_name, student_id
        )
        os.makedirs(student_dir, exist_ok=True)

        existing_submissions = [
            d for d in os.listdir(student_dir) if d.startswith("submission-")
        ]
        version = len(existing_submissions) + 1
        submission_path = os.path.join(student_dir, f"submission-{version}")
        os.makedirs(submission_path, exist_ok=True)

        filename = secure_filename(file.filename)
        filepath = os.path.join(submission_path, filename)
        file.save(filepath)

        try:
            if filepath.endswith(".zip"):
                with zipfile.ZipFile(filepath, "r") as zip_ref:
                    zip_ref.extractall(submission_path)
            else:
                with tarfile.open(filepath, "r:*") as tar:
                    tar.extractall(path=submission_path)
        except (tarfile.ReadError, zipfile.BadZipFile):
            return jsonify({"error": "Failed to extract archive."}), 400

        # Start the container execution and logging in a background thread
        thread = threading.Thread(
            target=run_and_log_result, args=(submission_path, homework_name, student_id)
        )
        thread.start()

        # Return a simple success page immediately
        return (
            f"""
        <h1>Submission Received!</h1>
        <p>Thank you, {student_id}. Your submission for <strong>{homework_name}</strong> (version {version}) has been received and is being processed.</p>
        <a href="/">Back to Home</a>
        """,
            202,
        )

    return jsonify({"error": "File type not allowed"}), 400


@app.route("/run", methods=["POST"])
def run_command():
    auth_token = request.headers.get("Authorization")
    if not auth_token or auth_token != f"Bearer {SECRET_TOKEN}":
        return jsonify({"error": "Unauthorized"}), 401

    data = request.get_json()
    if not data or "command" not in data:
        return jsonify({"error": "No command provided"}), 400

    command = data["command"]
    student_id = data.get("student_id")
    homework_name = data.get("homework_name")  # New: specify homework

    run_dir = None
    temp_dir = None

    try:
        if student_id and homework_name:
            # Path is now submissions/<homework_name>/<student_id>
            student_dir = os.path.join(
                app.config["UPLOAD_FOLDER"], homework_name, student_id
            )
            if not os.path.isdir(student_dir):
                return (
                    jsonify(
                        {
                            "error": f"No submissions found for student {student_id} in {homework_name}"
                        }
                    ),
                    404,
                )

            submissions = [
                d for d in os.listdir(student_dir) if d.startswith("submission-")
            ]
            if not submissions:
                return (
                    jsonify(
                        {
                            "error": f"No submissions found for student {student_id} in {homework_name}"
                        }
                    ),
                    404,
                )

            # Find the latest submission
            latest_submission = max(submissions, key=lambda s: int(s.split("-")[1]))
            run_dir = os.path.join(student_dir, latest_submission)
            working_dir = "/submission"
            volumes = {os.path.abspath(run_dir): {"bind": working_dir, "mode": "rw"}}
        else:
            # Run in an empty temporary directory if no student/homework is specified
            temp_dir = tempfile.TemporaryDirectory()
            run_dir = temp_dir.name
            working_dir = "/workspace"
            volumes = {os.path.abspath(run_dir): {"bind": working_dir, "mode": "rw"}}

        container = client.containers.run(
            "homework-runner",
            command=command,
            volumes=volumes,
            working_dir=working_dir,
            detach=True,
        )
        result = container.wait()
        logs = container.logs().decode("utf-8")
        container.remove()

        return jsonify(
            {
                "student_id": student_id or "N/A",
                "homework": homework_name or "N/A",
                "status_code": result["StatusCode"],
                "logs": logs,
            }
        )
    except docker.errors.ImageNotFound:
        return (
            jsonify(
                {
                    "error": 'Docker image "homework-runner" not found. Please build it first.'
                }
            ),
            500,
        )
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    finally:
        if temp_dir:
            temp_dir.cleanup()


# --- Static Site Serving ---


@app.route("/")
def serve_index():
    return send_from_directory(SITE_DIR, "index.html")


@app.route("/<path:path>")
def serve_static_files(path):
    return send_from_directory(SITE_DIR, path)


def copy_run_py():
    source_file = "run.py"
    destination_dir = "/tmp"
    destination_file = os.path.join(destination_dir, "run.py")

    if not os.path.exists(source_file):
        print(f"Error: Source file '{source_file}' not found in the current directory.")
    else:
        try:
            os.makedirs(destination_dir, exist_ok=True)
            shutil.copy(source_file, destination_file)
            current_permissions = os.stat(destination_file).st_mode
            new_permissions = (
                current_permissions | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
            )
            os.chmod(destination_file, new_permissions)
        except Exception as e:
            print(f"An error occurred: {e}")


if __name__ == "__main__":
    # Build the static site first
    if os.system("mkdocs build") != 0:
        print("Failed to build MkDocs site.")
        exit(1)

    copy_run_py()
    os.makedirs(UPLOAD_FOLDER, exist_ok=True)
    app.run(threaded=False, debug=True, port=8080)