from flask import Flask, request, render_template_string, send_file, abort, redirect, url_for, Response
import os
import subprocess
import tempfile
from datetime import datetime
import shutil
import json
import uuid
from functools import wraps

app = Flask(__name__)

# Paths to dynamic storage
TOKEN_DB_PATH = "/opt/wgconfig/tokens.json"
SUBNET_DB_PATH = "/opt/wgconfig/subnets.json"
USED_LOG_PATH = "/opt/wgconfig/used_tokens.log"

SERVER_PUBLIC_KEY = "bsped9nemtP9fRfcVLoLL3JWTj7zHFYWJebAhgH7RFw="
SERVER_ENDPOINT = "24.249.187.133:51820"
PEERS_CONF_PATH = "/opt/wgconfig/wg0-peers.conf"
KEY_LOG_PATH = "/var/log/wg_keys_issued.log"

ADMIN_USERNAME = "admin"
ADMIN_PASSWORD = ""

HTML_FORM = '''
<!doctype html>
<title>WireGuard Config Generator</title>
<h2>WireGuard Config Generator</h2>
<p>Welcome, your token is valid. Click below to generate your secure config.</p>
<form method="post">
    <button type="submit">Generate My WireGuard Config</button>
</form>
<p><a href="/instructions" target="_blank">How to use this page</a></p>
'''

HTML_INSTRUCTIONS = '''
<!doctype html>
<title>How to Use This Page</title>
<h2>📜 How to Use This Page: WireGuard Config Generator</h2>
<p>Welcome! This page allows you to generate a secure, one-time WireGuard VPN configuration. Follow these instructions carefully:</p>
<h3>✅ Step 1: Access the Page with Your Token</h3>
<ol>
<li>You must have a <strong>valid access token</strong> (provided by the administrator).</li>
<li>Open your browser and go to:<br>
<code>https://wgsetup.glorytoyah.org/setup?token=YourToken</code></li>
<li>Each token is valid for <strong>one-time use</strong>.</li>
</ol>
<h3>✅ Step 2: Generate Your Configuration</h3>
<ol>
<li>If your token is valid, you’ll see a <strong>“Generate My WireGuard Config”</strong> button.</li>
<li>Click the button.</li>
<li>A <code>.conf</code> file will be downloaded — this is your <strong>personal VPN configuration file</strong>.</li>
</ol>
<h3>⚠️ Important Notes:</h3>
<ul>
<li>Your configuration includes a <strong>private key</strong>. Keep this file safe — it should <strong>never be shared</strong>.</li>
<li>Each token is tied to a specific VPN IP address. <strong>Do not reuse the token for multiple devices.</strong></li>
</ul>
<h3>✅ Step 3: Import Config into WireGuard</h3>
<h4>📱 On Mobile (iOS/Android):</h4>
<ol>
<li>Download the <strong>WireGuard app</strong> from the App Store or Google Play.</li>
<li>Open the app and tap <strong>Add</strong> → <strong>Import from file or archive</strong>.</li>
<li>Locate and select the downloaded <code>.conf</code> file.</li>
</ol>
<h4>💻 On Desktop (Windows/macOS/Linux):</h4>
<ol>
<li>Install WireGuard from <a href="https://www.wireguard.com/install/" target="_blank">wireguard.com</a></li>
<li>Open the app and click <strong>Add Tunnel</strong> → <strong>Import from file</strong>.</li>
<li>Browse to and import your <code>.conf</code> file.</li>
</ol>
<h3>🟢 Once Imported</h3>
<p>Activate the VPN tunnel. You are now securely connected to your private network.</p>
'''

def check_auth(username, password):
    return username == ADMIN_USERNAME and password == ADMIN_PASSWORD

def authenticate():
    return Response("Login Required", 401, {"WWW-Authenticate": "Basic realm='Login Required'"})

def requires_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth or not check_auth(auth.username, auth.password):
            return authenticate()
        return f(*args, **kwargs)
    return decorated

def load_tokens():
    if os.path.exists(TOKEN_DB_PATH):
        with open(TOKEN_DB_PATH, "r") as f:
            tokens = json.load(f)
            upgraded = {}
            for k, v in tokens.items():
                if isinstance(v, str):
                    upgraded[k] = {"ip": v, "group": "unknown"}
                else:
                    upgraded[k] = v
            return upgraded
    return {}

def save_tokens(tokens):
    with open(TOKEN_DB_PATH, "w") as f:
        json.dump(tokens, f, indent=2)

def load_subnets():
    if os.path.exists(SUBNET_DB_PATH):
        with open(SUBNET_DB_PATH, "r") as f:
            return json.load(f)
    return {}

def save_subnets(subnets):
    with open(SUBNET_DB_PATH, "w") as f:
        json.dump(subnets, f, indent=2)

@app.route("/admin/subnets", methods=["GET", "POST"])
@requires_auth
def manage_subnets():
    subnets = load_subnets()
    message = ""
    if request.method == "POST":
        name = request.form.get("name")
        base = request.form.get("base")
        if name and base:
            subnets[name] = base
            save_subnets(subnets)
            message = f"Added subnet: {name} - {base}"
    html = "<h2>Subnet Groups</h2><table border='1'><tr><th>Group</th><th>Base Subnet</th></tr>"
    for name, base in subnets.items():
        html += f"<tr><td>{name}</td><td>{base}</td></tr>"
    html += "</table><h3>Add New Subnet</h3>"
    html += f"<p style='color:green'>{message}</p>" if message else ""
    html += """
    <form method='post'>
        Group Name: <input name='name' required>
        Base Subnet (e.g. 10.0.10): <input name='base' required>
        <button type='submit'>Add Subnet</button>
    </form>
    <p><a href='/admin/tokens'>Back to Token Dashboard</a></p>
    """
    return html

@app.route("/generate-and-redirect/<group>")
def generate_and_redirect(group):
    tokens = load_tokens()
    subnets = load_subnets()
    if group not in subnets:
        return "Invalid group name", 400
    base = subnets[group]
    used_ips = [data["ip"] for data in tokens.values() if data.get("group") == group and data.get("ip", "").startswith(base)]
    next_host = 2 + len(used_ips)
    assigned_ip = f"{base}.{next_host}/24"
    token = str(uuid.uuid4())
    tokens[token] = {"ip": assigned_ip, "group": group}
    save_tokens(tokens)
    return redirect(url_for("setup", token=token))

@app.route("/admin/tokens")
@requires_auth
def view_tokens():
    tokens = load_tokens()
    html = "<h2>Active Tokens</h2><table border='1'><tr><th>Token</th><th>IP</th><th>Group</th><th>Action</th></tr>"
    for t, v in tokens.items():
        html += f"<tr><td>{t}</td><td>{v['ip']}</td><td>{v['group']}</td><td><a href='/admin/revoke/{t}'>Revoke</a></td></tr>"
    html += "</table><p><a href='/admin/used'>View Used Tokens</a> | <a href='/admin/subnets'>Manage Subnets</a></p>"
    return html

@app.route("/admin/revoke/<token>")
@requires_auth
def revoke_token(token):
    tokens = load_tokens()
    if token in tokens:
        with open(USED_LOG_PATH, "a") as f:
            f.write(f"[REVOKED] {token} - {tokens[token]}\n")
        tokens.pop(token)
        save_tokens(tokens)
    return redirect(url_for("view_tokens"))

@app.route("/admin/used")
@requires_auth
def view_used():
    if not os.path.exists(USED_LOG_PATH):
        return "No used/revoked tokens logged."
    with open(USED_LOG_PATH) as f:
        return f"<h2>Used or Revoked Tokens</h2><pre>{f.read()}</pre>"

@app.route("/setup", methods=["GET", "POST"])
def setup():
    token = request.args.get("token")
    tokens = load_tokens()

    if token not in tokens:
        return abort(403, description="Invalid or missing token.")

    assigned_ip = tokens[token]["ip"]

    if request.method == "POST":
        try:
            wg_path = shutil.which("wg")
            if not wg_path:
                raise FileNotFoundError("WireGuard 'wg' binary not found in system path.")
            private_key = subprocess.check_output(["wg", "genkey"]).decode().strip()
            public_key = subprocess.check_output(["wg", "pubkey"], input=private_key.encode()).decode().strip()
        except Exception as e:
            app.logger.error(f"Failed to run WireGuard key generation: {e}")
            return "Internal error: key generation failed", 500
        config_content = f"""
[Interface]
PrivateKey = {private_key}
Address = {assigned_ip}

[Peer]
PublicKey = {SERVER_PUBLIC_KEY}
Endpoint = {SERVER_ENDPOINT}
AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 25
"""
        temp_dir = tempfile.mkdtemp()
        config_path = os.path.join(temp_dir, f"wgclient_{token}.conf")
        with open(config_path, "w") as f:
            f.write(config_content)
        try:
            with open(KEY_LOG_PATH, "a", encoding="utf-8") as log:
                log.write(f"[{datetime.now()}] {token} -> {public_key} ({assigned_ip})\n")
        except Exception as e:
            app.logger.error(f"Failed to write to key log: {e}")
        peer_entry = f"\n[Peer]\nPublicKey = {public_key}\nAllowedIPs = {assigned_ip.split('/')[0]}/32\n"
        try:
            os.makedirs(os.path.dirname(PEERS_CONF_PATH), exist_ok=True)
            with open(PEERS_CONF_PATH, "a") as peerconf:
                peerconf.write(peer_entry)
            subprocess.run(["sudo", "wg", "addconf", "wg0", PEERS_CONF_PATH], check=True)
        except Exception as e:
            app.logger.error(f"Failed to update live WireGuard config: {e}")
        with open(USED_LOG_PATH, "a") as log:
            log.write(f"[USED] {token} - {tokens[token]}\n")
        tokens.pop(token)
        save_tokens(tokens)
        return send_file(config_path, as_attachment=True)
    return render_template_string(HTML_FORM)

@app.route("/instructions")
def instructions():
    return render_template_string(HTML_INSTRUCTIONS)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)
