import base64, io, os, zipfile
import gradio as gr
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey

def _wg_key(pk): return base64.b64encode(pk.private_bytes_raw()).decode()
def _wg_pub(pk): return base64.b64encode(pk.public_key().public_bytes_raw()).decode()

def generate(ip_a, ip_b, wg_port, psk_enabled):
    priv_a, priv_b, priv_c = (X25519PrivateKey.generate() for _ in range(3))
    ka_priv, ka_pub = _wg_key(priv_a), _wg_pub(priv_a)
    kb_priv, kb_pub = _wg_key(priv_b), _wg_pub(priv_b)
    kc_priv, kc_pub = _wg_key(priv_c), _wg_pub(priv_c)

    psk_line = ""
    if psk_enabled:
        import secrets
        psk = base64.b64encode(secrets.token_bytes(32)).decode()
        psk_line = f"PresharedKey = {psk}\n"

    node_a = f"""[Interface]
# Node A — {ip_a}
Address = 10.0.0.1/24
ListenPort = {int(wg_port)}
PrivateKey = {ka_priv}

PostUp = sysctl -w net.ipv4.ip_forward=1
PostUp = iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE
PostDown = iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE

[Peer]
# Node B
PublicKey = {kb_pub}
Endpoint = {ip_b}:{int(wg_port)}
{psk_line}AllowedIPs = 10.0.0.2/32
PersistentKeepalive = 25

[Peer]
# Client
PublicKey = {kc_pub}
{psk_line}AllowedIPs = 10.0.0.3/32
PersistentKeepalive = 25
"""

    node_b = f"""[Interface]
# Node B — {ip_b}
Address = 10.0.0.2/24
ListenPort = {int(wg_port)}
PrivateKey = {kb_priv}

PostUp = sysctl -w net.ipv4.ip_forward=1
PostUp = iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE
PostDown = iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE

[Peer]
# Node A
PublicKey = {ka_pub}
Endpoint = {ip_a}:{int(wg_port)}
{psk_line}AllowedIPs = 10.0.0.1/32, 10.0.0.3/32
PersistentKeepalive = 25
"""

    client = f"""[Interface]
# Client
Address = 10.0.0.3/24
PrivateKey = {kc_priv}
DNS = 1.1.1.1, 8.8.8.8

[Peer]
# Gateway — Node A
PublicKey = {ka_pub}
Endpoint = {ip_a}:{int(wg_port)}
{psk_line}AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 25
"""

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as z:
        z.writestr("node_a_wg0.conf", node_a)
        z.writestr("node_b_wg0.conf", node_b)
        z.writestr("client_wg0.conf", client)
    buf.seek(0)
    out = "/tmp/gateway_configs.zip"
    with open(out, "wb") as f: f.write(buf.read())

    preview = f"=== Node A ===\n{node_a}\n=== Node B ===\n{node_b}\n=== Client ===\n{client}"
    return preview, out

CSS = "body{background:#0a0a0f}.gradio-container{max-width:900px!important}footer{display:none!important}"

with gr.Blocks(title="Gateway Router", css=CSS,
    theme=gr.themes.Base(primary_hue="blue", neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "sans-serif"])) as demo:

    gr.HTML("""<div style='background:linear-gradient(135deg,#0d1117,#161b22);
        border-bottom:1px solid #1e3a5f;padding:14px 24px;margin-bottom:8px'>
        <div style='font-size:1.3rem;font-weight:700;letter-spacing:.08em;
            background:linear-gradient(90deg,#58a6ff,#7ee8fa);
            -webkit-background-clip:text;-webkit-text-fill-color:transparent'>
            GATEWAY ROUTER</div>
        <div style='font-size:.78rem;color:#8b949e'>
            Dual-node WireGuard VPN · Pure Python key generation · No installation required
        </div></div>""")

    with gr.Row():
        with gr.Column():
            ip_a    = gr.Textbox(label="Node A Public IP", placeholder="1.2.3.4")
            ip_b    = gr.Textbox(label="Node B Public IP", placeholder="5.6.7.8")
            wg_port = gr.Number(label="WireGuard Port", value=51820, precision=0)
            psk     = gr.Checkbox(label="Add Preshared Key (recommended)", value=True)
            btn     = gr.Button("Generate Configs ▶", variant="primary")
        with gr.Column():
            preview = gr.Textbox(label="Config Preview", lines=20, interactive=False)
            dl      = gr.File(label="Download ZIP")

    btn.click(fn=generate, inputs=[ip_a, ip_b, wg_port, psk], outputs=[preview, dl])

demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))