from flask import Flask, request, make_response
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import os

app = Flask(__name__)

KEY = os.urandom(16)
registered = set()

def encrypt(user_id: int) -> str:
    IV = os.urandom(16)
    cipher = AES.new(KEY, AES.MODE_CBC, IV)
    pt = pad(str(user_id).encode(), 16)
    ct = cipher.encrypt(pt)
    return ct.hex() + IV.hex()

def decrypt(cookie: str) -> int:
    ct = bytes.fromhex(cookie)
    IV = ct[16:]
    ct = ct[:16]
    cipher = AES.new(KEY, AES.MODE_CBC, IV)
    pt = unpad(cipher.decrypt(ct), 16)
    return int(pt.decode())

@app.route("/register")
def register():
    user_id = 2  # admin = 1
    session = encrypt(user_id)
    resp = make_response("Registered!")
    resp.set_cookie("session", session)
    return resp

@app.route("/flag")
def flag():
    session = request.cookies.get("session")
    try:
        user_id = decrypt(session)
    except:
        return "Invalid session."
    if user_id == 1:
        return "HZU18{flag was here xD.}"
    return "You are not admin."

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