import os
import base64
import hashlib
import gradio as gr
from pathlib import Path
from Crypto.Cipher import AES
from Crypto.Util import Counter

ROOT = Path(__file__).parent
ASSET_IMG = ROOT / "mascot.png"
REQUIREMENTS = ROOT / "requirements.txt"

def _key_and_nonce():
    # Key = SHA256 of mascot.png bytes
    img_bytes = (ASSET_IMG.read_bytes() if ASSET_IMG.exists() else b"")
    key = hashlib.sha256(img_bytes).digest()  # 32 bytes (AES-256)

    # Nonce (8 bytes) derived from first line of requirements.txt
    rl = ""
    if REQUIREMENTS.exists():
        with REQUIREMENTS.open("r", encoding="utf-8", errors="ignore") as f:
            rl = (f.readline() or "").strip()
    nonce_src = hashlib.sha1(rl.encode("utf-8")).digest()
    nonce = nonce_src[:8] if nonce_src else b"\x00" * 8

    return key, nonce

def _encrypt_flag(flag: str) -> str:
    key, nonce = _key_and_nonce()
    ctr = Counter.new(64, prefix=nonce, initial_value=0)
    cipher = AES.new(key, AES.MODE_CTR, counter=ctr)
    ct = cipher.encrypt(flag.encode("utf-8"))
    return base64.b64encode(ct).decode("ascii")

def get_ciphertext():
    flag = (os.environ.get("FLAG") or "").strip()
    if not flag:
        return "⚠️ FLAG not set."
    return _encrypt_flag(flag)

def verify(candidate: str):
    flag = (os.environ.get("FLAG") or "").strip()
    if not flag:
        return ("⚠️ FLAG not set.", "")
    candidate = (candidate or "").strip()
    if not candidate:
        return ("No input provided.", "")
    if candidate == flag:
        # Always return the official flag from env (not user input)
        return ("✅ Correct.", flag)
    return ("❌ Incorrect.", "")

custom_css = """
#hero { display: flex; gap: 16px; align-items: center; }
#hero .title h1 { margin: 0; }
"""

with gr.Blocks(title="ENUSEC Freshers – Decrypter 2025", css=custom_css) as demo:
    # Hero row with mascot + title
    with gr.Row(elem_id="hero"):
        gr.Image(
            value=str(ASSET_IMG) if ASSET_IMG.exists() else None,
            show_label=False,
            interactive=False,
            height=160,
            visible=ASSET_IMG.exists()
        )
        gr.Markdown(
            """
# Decrypter 2025
Click **Get Encrypted Blob** to receive a Base64 string.  
Submit the flag when you have it.
            """,
            elem_classes=["title"]
        )

    # Ciphertext section
    b = gr.Button("Get Encrypted Blob")
    blob = gr.Textbox(label="Ciphertext (Base64)", interactive=False)
    b.click(get_ciphertext, None, blob)

    # Submission section
    gr.Markdown("## Submit Flag")
    guess = gr.Textbox(label="freshers{...}")
    submit = gr.Button("Submit")
    msg = gr.Markdown()
    out = gr.Textbox(label="Official Flag", interactive=False)
    submit.click(verify, guess, [msg, out])

if __name__ == "__main__":
    demo.launch()