from flask import Flask, request, jsonify
import base64
import subprocess
import binascii

from db import get_enc_flag_and_alg, get_flag_and_tmp_flag, update_data
from crypto import base_enc, xor_enc, rotr_enc, rotl_enc


def encrypt_tmp_flag(tmp_flag, alg):
    if alg.startswith("base:"):
        base_number = int(alg.split(":")[1])
        return base_enc(tmp_flag, base_number)
    if alg == "xor":
        return xor_enc(tmp_flag)
    if alg == "rotr":
        bits = len(tmp_flag.encode("utf-8")) * 8
        return rotr_enc(tmp_flag, bits)
    if alg == "rotl":
        bits = len(tmp_flag.encode("utf-8")) * 8
        return rotl_enc(tmp_flag, bits)
    raise ValueError("Неизвестный алгоритм")


app = Flask(__name__)
app.config["MAX_CONTENT_LENGTH"] = 4096


@app.route('/get_question', methods=['GET'])
def generate_question():
    data = request.get_json(silent=True)
    if not isinstance(data, dict) or "round_id" not in data:
        return "Некорректный запрос", 400
    id = data["round_id"]
    result = get_enc_flag_and_alg(id)
    if result is None:
        return "Такой вопрос не найден", 404
    enc_flag, alg = result
    question = (f"Чтобы получить флаг, вы, должны решить ЛёГкУю задачку. Ваш билет таков.\nРашифруйте: {enc_flag}\nАлгоритм: {alg}")
    return question


@app.route('/check_answer', methods=['POST'])
def check():
    data = request.get_json(silent=True)
    if not isinstance(data, dict):
        return "Некорректный запрос", 400
    id = data.get("round_id")
    enc_ans = data.get("answer")
    if not isinstance(id, str) or not isinstance(enc_ans, str):
        return "Некорректный запрос", 400
    try:
        dec_ans = base64.b64decode(enc_ans, validate = True)
    except (binascii.Error, ValueError):
        return "Некорректный Base64", 400
    if len(dec_ans) > 1024:
        return "Слишком длинный ответ", 413
    result = get_flag_and_tmp_flag(id)
    if result is None:
        return "Такой вопрос не найден", 404
    flag, answer = result
    try:
        res = subprocess.run(["./validator/validate", flag, answer], input = dec_ans, capture_output = True, timeout = 2)
    except subprocess.TimeoutExpired:
        return "Validator timeout", 504
    return res.stdout


@app.route('/put_question', methods=['POST'])
def put_question():
    data = request.get_json(silent=True)
    if not isinstance(data, dict):
        return jsonify({"ok": False, "error": "Ожидался JSON"}), 400
    round_id = data.get("round_id")
    flag = data.get("flag")
    tmp_flag = data.get("tmp_flag")
    alg = data.get("alg")
    if not all(isinstance(value, str) for value in (round_id, flag, tmp_flag, alg)):
        return jsonify({"ok": False, "error": "Некорректные поля"}), 400
    if len(round_id) > 128:
        return jsonify({"ok": False, "error": "Слишком длинный round_id"}), 400
    if len(flag) > 256 or len(tmp_flag) > 256:
        return jsonify({"ok": False, "error": "Слишком длинный флаг"}), 400
    if len(alg) > 32:
        return jsonify({"ok": False, "error": "Слишком длинное название алгоритма"}), 400
    try:
        enc_flag = encrypt_tmp_flag(tmp_flag, alg)
    except (ValueError, TypeError):
        return jsonify({"ok": False, "error": "Ошибка при шифровании"}), 400
    result = update_data(round_id, flag, tmp_flag, enc_flag, alg)
    if result is False:
        return jsonify({"ok": False, "error": "Такой round_id уже существует"}), 409
    return jsonify({"ok": True})


if __name__ == "__main__":
    app.run(host = "0.0.0.0", port = 1337, debug = False)