#!/usr/bin/env python3
"""
dsnat-setup.py — DSNAT Wireguard Tunnel Setup
Aya's Tengu Express Edition 🗞️
"""

import curses
import subprocess
import os
import sys
import time
import textwrap
import re
import tempfile
import shutil
from pathlib import Path

# ─── Palette ──────────────────────────────────────────────────────────────────
C_TITLE     = 1
C_BORDER    = 2
C_NORMAL    = 3
C_HIGHLIGHT = 4
C_ERROR     = 5
C_SUCCESS   = 6
C_DIM       = 7
C_INPUT     = 8

def init_colors():
    curses.start_color()
    curses.use_default_colors()
    curses.init_pair(C_TITLE,     curses.COLOR_CYAN,    -1)
    curses.init_pair(C_BORDER,    curses.COLOR_BLUE,    -1)
    curses.init_pair(C_NORMAL,    curses.COLOR_WHITE,   -1)
    curses.init_pair(C_HIGHLIGHT, curses.COLOR_BLACK,   curses.COLOR_CYAN)
    curses.init_pair(C_ERROR,     curses.COLOR_RED,     -1)
    curses.init_pair(C_SUCCESS,   curses.COLOR_GREEN,   -1)
    curses.init_pair(C_DIM,       curses.COLOR_WHITE,   -1)
    curses.init_pair(C_INPUT,     curses.COLOR_YELLOW,  -1)

# ─── Drawing helpers ──────────────────────────────────────────────────────────

def draw_box(win, y, x, h, w, color=C_BORDER, title=None):
    """Draw a single-line box, optionally with a title.
    Fully bounds-safe: clamps to terminal size and swallows addch ERR on the
    bottom-right corner cell (curses raises there on most terminals)."""
    max_y, max_x = win.getmaxyx()
    if y < 0 or x < 0 or y >= max_y or x >= max_x:
        return
    h = min(h, max_y - y)
    w = min(w, max_x - x)
    if h < 2 or w < 2:
        return

    attr = curses.color_pair(color)

    def _ch(cy, cx, ch):
        if 0 <= cy < max_y and 0 <= cx < max_x:
            try:
                win.addch(cy, cx, ch, attr)
            except curses.error:
                pass

    _ch(y,     x,     curses.ACS_ULCORNER)
    _ch(y,     x+w-1, curses.ACS_URCORNER)
    _ch(y+h-1, x,     curses.ACS_LLCORNER)
    _ch(y+h-1, x+w-1, curses.ACS_LRCORNER)
    for i in range(1, w-1):
        _ch(y,     x+i, curses.ACS_HLINE)
        _ch(y+h-1, x+i, curses.ACS_HLINE)
    for j in range(1, h-1):
        _ch(y+j, x,     curses.ACS_VLINE)
        _ch(y+j, x+w-1, curses.ACS_VLINE)

    if title:
        label = f" {title} "
        tx = x + (w - len(label)) // 2
        if 0 <= y < max_y and tx >= 0:
            try:
                win.addstr(y, tx, label[:max_x - tx],
                           curses.color_pair(C_TITLE) | curses.A_BOLD)
            except curses.error:
                pass

def safe_addstr(win, y, x, s, attr=0):
    max_y, max_x = win.getmaxyx()
    if y < 0 or y >= max_y or x < 0 or x >= max_x:
        return
    available = max_x - x - 1
    if available <= 0:
        return
    try:
        win.addstr(y, x, s[:available], attr)
    except curses.error:
        pass

def center_str(win, y, s, attr=0):
    _, w = win.getmaxyx()
    x = max(0, (w - len(s)) // 2)
    safe_addstr(win, y, x, s, attr)

# ─── Banner ───────────────────────────────────────────────────────────────────

BANNER = [
    r"  ____  ____  _   _    _  _____   ",
    r" |  _ \/ ___|| \ | |  / \|_   _|  ",
    r" | | | \___ \|  \| | / _ \ | |    ",
    r" | |_| |___) | |\  |/ ___ \| |    ",
    r" |____/|____/|_| \_/_/   \_\_|    ",
    r"  Wireguard Tunnel Setup  v1.0     ",
]

def draw_banner(win, start_y=1):
    for i, line in enumerate(BANNER):
        center_str(win, start_y + i, line,
                   curses.color_pair(C_TITLE) | curses.A_BOLD)

# ─── Role Selection ───────────────────────────────────────────────────────────

def select_role(stdscr):
    curses.curs_set(0)
    h, w = stdscr.getmaxyx()
    options = ["  Edge (VPS)  ", "  Client (Home server)  "]
    selected = 0

    while True:
        stdscr.erase()
        draw_banner(stdscr, 1)

        center_str(stdscr, 9,
            "Aya's DSNAT Tunnel Installer 🗞️",
            curses.color_pair(C_NORMAL) | curses.A_BOLD)
        center_str(stdscr, 10,
            "Select your role:",
            curses.color_pair(C_DIM))

        box_w = 40
        box_x = (w - box_w) // 2
        draw_box(stdscr, 12, box_x, 7, box_w, C_BORDER, "Role")

        for i, opt in enumerate(options):
            attr = curses.color_pair(C_HIGHLIGHT) | curses.A_BOLD if i == selected else curses.color_pair(C_NORMAL)
            label = f" {'>' if i == selected else ' '} {opt}"
            safe_addstr(stdscr, 14 + i*2, box_x + 2, label.ljust(box_w - 4), attr)

        center_str(stdscr, 20,
            "↑/↓ Navigate   Enter Select   Q Quit",
            curses.color_pair(C_DIM))

        stdscr.refresh()
        key = stdscr.getch()
        if key in (curses.KEY_UP, ord('k')) and selected > 0:
            selected -= 1
        elif key in (curses.KEY_DOWN, ord('j')) and selected < len(options) - 1:
            selected += 1
        elif key in (curses.KEY_ENTER, 10, 13):
            return "edge" if selected == 0 else "client"
        elif key in (ord('q'), ord('Q')):
            return None

# ─── Input form ───────────────────────────────────────────────────────────────

def input_box(stdscr, prompt, default="", secret=False):
    """Single-line input with a prompt. Returns the entered string."""
    curses.curs_set(1)
    h, w = stdscr.getmaxyx()
    box_w = min(70, w - 4)
    box_x = (w - box_w) // 2
    box_y = h // 2 - 3

    draw_box(stdscr, box_y, box_x, 5, box_w, C_BORDER, "Input")
    wrapped = textwrap.wrap(prompt, box_w - 4)
    for i, line in enumerate(wrapped[:2]):
        safe_addstr(stdscr, box_y + 1 + i, box_x + 2, line,
                    curses.color_pair(C_NORMAL))

    field_y = box_y + 3
    field_x = box_x + 2
    field_w = box_w - 4

    buf = list(default)
    cursor = len(buf)

    while True:
        display = "".join(buf)
        if secret:
            display = "*" * len(buf)
        padded = (display[max(0, cursor - field_w + 1):] if cursor >= field_w
                  else display)
        safe_addstr(stdscr, field_y, field_x,
                    padded.ljust(field_w)[:field_w],
                    curses.color_pair(C_INPUT) | curses.A_UNDERLINE)
        vis_x = field_x + min(cursor, field_w - 1)
        stdscr.move(field_y, vis_x)
        stdscr.refresh()

        key = stdscr.getch()
        if key in (curses.KEY_ENTER, 10, 13):
            curses.curs_set(0)
            return "".join(buf)
        elif key in (curses.KEY_BACKSPACE, 127, 8):
            if cursor > 0:
                buf.pop(cursor - 1)
                cursor -= 1
        elif key == curses.KEY_DC:
            if cursor < len(buf):
                buf.pop(cursor)
        elif key == curses.KEY_LEFT and cursor > 0:
            cursor -= 1
        elif key == curses.KEY_RIGHT and cursor < len(buf):
            cursor += 1
        elif key == curses.KEY_HOME:
            cursor = 0
        elif key == curses.KEY_END:
            cursor = len(buf)
        elif 32 <= key <= 126:
            buf.insert(cursor, chr(key))
            cursor += 1

def collect_edge_params(stdscr):
    h, w = stdscr.getmaxyx()
    params = {}

    fields = [
        ("wg_port",    "WireGuard listen port:",              "51820"),
        ("iface",      "Public network interface (e.g. ens3 or eth0):", "ens3"),
        ("exposed_ports", "Ports to DNAT/expose (comma-sep, e.g. 80,443,22):", "80,443,22"),
        ("ssh_port",   "SSH port for sshd_config (shame-edges standard):", "42"),
        ("client_pub", "Client public key (leave blank to generate now):", ""),
    ]

    for key, prompt, default in fields:
        stdscr.erase()
        draw_banner(stdscr, 1)
        center_str(stdscr, 9, "Edge Configuration", curses.color_pair(C_TITLE) | curses.A_BOLD)
        val = input_box(stdscr, prompt, default)
        params[key] = val if val.strip() else default

    return params

def collect_client_params(stdscr):
    params = {}

    fields = [
        ("vps_ip",    "VPS public IP address:", ""),
        ("wg_port",   "Edge WireGuard port:",   "51820"),
        ("edge_pub",  "Edge public key:",        ""),
    ]

    for key, prompt, default in fields:
        stdscr.erase()
        draw_banner(stdscr, 1)
        center_str(stdscr, 9, "Client Configuration", curses.color_pair(C_TITLE) | curses.A_BOLD)
        val = input_box(stdscr, prompt, default)
        params[key] = val if val.strip() else default

    return params

# ─── Step runner ──────────────────────────────────────────────────────────────

class StepRunner:
    """Runs a list of (label, callable) steps with a live TUI log."""

    def __init__(self, stdscr, steps, title="Running…"):
        self.stdscr = stdscr
        self.steps  = steps
        self.title  = title
        self.log    = []
        self.errors = []

    def _redraw(self, current_idx, status="running"):
        s = self.stdscr
        h, w = s.getmaxyx()
        s.erase()
        draw_banner(s, 1)
        center_str(s, 9, self.title,
                   curses.color_pair(C_TITLE) | curses.A_BOLD)

        # Steps list
        box_w = min(70, w - 4)
        box_x = (w - box_w) // 2
        n = len(self.steps)
        box_h = n + 4
        draw_box(s, 11, box_x, box_h, box_w, C_BORDER, "Steps")

        for i, (label, _) in enumerate(self.steps):
            if i < current_idx:
                icon = "✓ "
                attr = curses.color_pair(C_SUCCESS) | curses.A_BOLD
            elif i == current_idx:
                if status == "error":
                    icon = "✗ "
                    attr = curses.color_pair(C_ERROR) | curses.A_BOLD
                else:
                    icon = "▶ "
                    attr = curses.color_pair(C_INPUT) | curses.A_BOLD
            else:
                icon = "  "
                attr = curses.color_pair(C_DIM)
            safe_addstr(s, 12 + i, box_x + 2, f"{icon}{label}"[:box_w - 4], attr)

        # Log tail
        log_y = 11 + box_h + 1
        log_h = h - log_y - 2
        if log_h > 3:
            draw_box(s, log_y, box_x, log_h, box_w, C_BORDER, "Log")
            visible = self.log[-(log_h - 2):]
            for i, line in enumerate(visible):
                safe_addstr(s, log_y + 1 + i, box_x + 2,
                            line[:box_w - 4],
                            curses.color_pair(C_DIM))

        s.refresh()

    def run(self):
        for idx, (label, fn) in enumerate(self.steps):
            self._redraw(idx, "running")
            try:
                fn(self._log)
            except Exception as e:
                self.errors.append(str(e))
                self._log(f"ERROR: {e}")
                self._redraw(idx, "error")
                self._pause("Step failed — press any key to continue…", C_ERROR)
                return False
        self._redraw(len(self.steps), "done")
        return True

    def _log(self, msg):
        for line in str(msg).splitlines():
            self.log.append(line)
        self._redraw_log_only()

    def _redraw_log_only(self):
        # light refresh just for the log pane
        s = self.stdscr
        h, w = s.getmaxyx()
        box_w = min(70, w - 4)
        box_x = (w - box_w) // 2
        n = len(self.steps)
        log_y = 11 + n + 4 + 1
        log_h = h - log_y - 2
        if log_h > 3:
            visible = self.log[-(log_h - 2):]
            for i, line in enumerate(visible):
                safe_addstr(s, log_y + 1 + i, box_x + 2,
                            " " * (box_w - 4), curses.color_pair(C_DIM))
                safe_addstr(s, log_y + 1 + i, box_x + 2,
                            line[:box_w - 4], curses.color_pair(C_DIM))
        s.refresh()

    def _pause(self, msg, color=C_NORMAL):
        h, w = self.stdscr.getmaxyx()
        center_str(self.stdscr, h - 2, msg, curses.color_pair(color))
        self.stdscr.refresh()
        self.stdscr.getch()

# ─── Shell helpers ────────────────────────────────────────────────────────────

def run(cmd, log_fn=None, check=True):
    if log_fn:
        log_fn(f"$ {cmd}")
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if log_fn:
        if result.stdout.strip():
            log_fn(result.stdout.strip())
        if result.stderr.strip():
            log_fn(result.stderr.strip())
    if check and result.returncode != 0:
        raise RuntimeError(f"Command failed (rc={result.returncode}): {cmd}\n{result.stderr}")
    return result

def wg_genkey():
    """Generate a WireGuard keypair using Python's cryptography library.
    No dependency on the `wg` binary — safe to call before apt install."""
    from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
    import base64
    priv = X25519PrivateKey.generate()
    priv_bytes = priv.private_bytes_raw()
    pub_bytes  = priv.public_key().public_bytes_raw()
    private = base64.b64encode(priv_bytes).decode()
    public  = base64.b64encode(pub_bytes).decode()
    return private, public

# ─── Config file generators ───────────────────────────────────────────────────

def make_wg_edge(private_key, client_pub, listen_port):
    return f"""[Interface]
Address = 10.44.0.1/24
PrivateKey = {private_key}
ListenPort = {listen_port}

[Peer]
PublicKey = {client_pub}
AllowedIPs = 10.44.0.2/32
"""

def make_wg_client(private_key, edge_pub, vps_ip, wg_port):
    return f"""[Interface]
Address = 10.44.0.2/24
PrivateKey = {private_key}

[Peer]
PublicKey = {edge_pub}
Endpoint = {vps_ip}:{wg_port}
AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 25
"""

def make_nftables(iface, exposed_ports, oci_mode):
    """
    Generate nftables.conf for the edge.
    oci_mode: if True, match on iifname instead of ip daddr.
    """
    ports_nft = "{" + ", ".join(p.strip() for p in exposed_ports.split(",")) + "}"

    if oci_mode:
        prerouting_rule = f'    iifname "{iface}" tcp dport {ports_nft} dnat to 10.44.0.2'
        comment = "# OCI mode: match iifname (public IP never appears on interface)"
    else:
        prerouting_rule = f'    tcp dport {ports_nft} dnat to 10.44.0.2'
        comment = "# Standard mode: DNAT all matching dst ports"

    return f"""#!/usr/sbin/nft -f
# DSNAT Tunnel — generated by dsnat-setup.py
{comment}

flush ruleset

table ip nat {{
  chain prerouting {{
    type nat hook prerouting priority -100;
{prerouting_rule}
  }}

  chain postrouting {{
    type nat hook postrouting priority 100;
    # Home server outbound traffic exits as VPS public IP
    ip saddr 10.44.0.2 oifname "{iface}" masquerade
  }}
}}

table ip filter {{

  set blacklist {{
    type ipv4_addr
    flags interval
  }}

  chain forward {{
    type filter hook forward priority 0; policy drop;
    ip saddr @blacklist drop
    # MSS clamp hardcoded to 1380 for WireGuard tunnel safety
    tcp flags syn tcp option maxseg size set 1380

    ct state established,related accept

    iifname "{iface}" oifname "wg0" accept
    iifname "wg0" oifname "{iface}" accept
  }}
}}
"""

def make_sysctl():
    return """net.ipv4.ip_forward=1
net.ipv4.conf.all.rp_filter=0
net.ipv4.conf.default.rp_filter=0
"""

def make_sysctl_oci(iface):
    """Extra rp_filter=0 for the specific OCI interface."""
    return make_sysctl() + f"net.ipv4.conf.{iface}.rp_filter=0\n"

# ─── OCI detection prompt ─────────────────────────────────────────────────────

def ask_oci(stdscr):
    h, w = stdscr.getmaxyx()
    options = ["  Yes, this is OCI (or similar hypervisor NAT)  ",
               "  No, standard VPS  "]
    selected = 0
    curses.curs_set(0)

    while True:
        stdscr.erase()
        draw_banner(stdscr, 1)
        center_str(stdscr, 9, "OCI / Hypervisor NAT Detection",
                   curses.color_pair(C_TITLE) | curses.A_BOLD)
        center_str(stdscr, 10,
            "Is this an OCI instance (or any VPS where the public IP is NOT on the interface)?",
            curses.color_pair(C_NORMAL))

        box_w = 58
        box_x = (w - box_w) // 2
        draw_box(stdscr, 12, box_x, 7, box_w, C_BORDER, "VPS Type")

        for i, opt in enumerate(options):
            attr = curses.color_pair(C_HIGHLIGHT) | curses.A_BOLD if i == selected else curses.color_pair(C_NORMAL)
            label = f" {'>' if i == selected else ' '} {opt}"
            safe_addstr(stdscr, 14 + i*2, box_x + 2, label.ljust(box_w - 4), attr)

        center_str(stdscr, 21, "↑/↓ Navigate   Enter Select",
                   curses.color_pair(C_DIM))
        stdscr.refresh()

        key = stdscr.getch()
        if key in (curses.KEY_UP, ord('k')) and selected > 0:
            selected -= 1
        elif key in (curses.KEY_DOWN, ord('j')) and selected < 1:
            selected += 1
        elif key in (curses.KEY_ENTER, 10, 13):
            return selected == 0

# ─── Summary / confirmation screen ───────────────────────────────────────────

def confirm_screen(stdscr, title, lines):
    """Show a scrollable summary. Returns True if user confirmed."""
    h, w = stdscr.getmaxyx()
    box_w = min(76, w - 4)
    box_x = (w - box_w) // 2
    BOX_Y = 10
    # bottom edge must be at h-3 (leaves room for the hint line at h-2)
    box_h = max(4, h - BOX_Y - 3)
    curses.curs_set(0)
    scroll = 0
    visible = box_h - 2

    while True:
        stdscr.erase()
        draw_banner(stdscr, 1)
        center_str(stdscr, 9, title,
                   curses.color_pair(C_TITLE) | curses.A_BOLD)
        draw_box(stdscr, BOX_Y, box_x, box_h, box_w, C_BORDER, "Review")

        for i in range(min(visible, len(lines) - scroll)):
            safe_addstr(stdscr, BOX_Y + 1 + i, box_x + 2,
                        lines[scroll + i][:box_w - 4],
                        curses.color_pair(C_NORMAL))

        center_str(stdscr, h - 2,
            "↑/↓ Scroll   Y Confirm & Install   N Abort",
            curses.color_pair(C_DIM))
        stdscr.refresh()

        key = stdscr.getch()
        if key in (curses.KEY_UP, ord('k')) and scroll > 0:
            scroll -= 1
        elif key in (curses.KEY_DOWN, ord('j')) and scroll < len(lines) - visible:
            scroll += 1
        elif key in (ord('y'), ord('Y')):
            return True
        elif key in (ord('n'), ord('N'), ord('q')):
            return False

# ─── Final info screen ────────────────────────────────────────────────────────

def info_screen(stdscr, title, lines, color=C_SUCCESS):
    h, w = stdscr.getmaxyx()
    box_w = min(76, w - 4)
    box_x = (w - box_w) // 2
    BOX_Y = 10
    box_h = max(4, h - BOX_Y - 3)
    curses.curs_set(0)
    scroll = 0
    visible = box_h - 2

    while True:
        stdscr.erase()
        draw_banner(stdscr, 1)
        center_str(stdscr, 9, title, curses.color_pair(color) | curses.A_BOLD)
        draw_box(stdscr, BOX_Y, box_x, box_h, box_w, C_BORDER)

        for i in range(min(visible, len(lines) - scroll)):
            safe_addstr(stdscr, BOX_Y + 1 + i, box_x + 2,
                        lines[scroll + i][:box_w - 4],
                        curses.color_pair(C_NORMAL))

        center_str(stdscr, h - 2, "↑/↓ Scroll   Q Quit", curses.color_pair(C_DIM))
        stdscr.refresh()

        key = stdscr.getch()
        if key in (curses.KEY_UP, ord('k')) and scroll > 0:
            scroll -= 1
        elif key in (curses.KEY_DOWN, ord('j')) and scroll < max(0, len(lines) - visible):
            scroll += 1
        elif key in (ord('q'), ord('Q')):
            return

# ─── Edge installer ───────────────────────────────────────────────────────────

def install_edge(stdscr):
    oci = ask_oci(stdscr)
    params = collect_edge_params(stdscr)

    iface      = params["iface"]
    wg_port    = params["wg_port"]
    ports      = params["exposed_ports"]
    ssh_port   = params["ssh_port"]
    client_pub_given = params["client_pub"].strip()

    # Generate keys (pure Python, no wg binary needed)
    edge_priv, edge_pub = wg_genkey()

    nft_conf   = make_nftables(iface, ports, oci)
    sysctl_conf = make_sysctl_oci(iface) if oci else make_sysctl()
    wg_conf    = make_wg_edge(edge_priv, client_pub_given or "<CLIENT_PUB_KEY>", wg_port)

    summary = [
        f"Role         : Edge (VPS)",
        f"OCI mode     : {'YES' if oci else 'no'}",
        f"Interface    : {iface}",
        f"WG port      : {wg_port}",
        f"Exposed ports: {ports}",
        f"SSH port     : {ssh_port}",
        f"Edge pub key : {edge_pub}",
        f"Client pub   : {client_pub_given or '(will need to fill in wg0.conf manually)'}",
        "",
        "── nftables.conf ──────────────────────────────────────",
        *nft_conf.splitlines(),
        "",
        "── /etc/wireguard/wg0.conf ────────────────────────────",
        *wg_conf.splitlines(),
        "",
        "── /etc/sysctl.d/99-forward.conf ──────────────────────",
        *sysctl_conf.splitlines(),
    ]

    if not confirm_screen(stdscr, "Edge — Review & Confirm", summary):
        return

    # Build steps
    def step_apt(log):
        run("apt-get update -qq", log)
        run("apt-get install -y wireguard nftables unattended-upgrades", log)
        run("dpkg-reconfigure -f noninteractive unattended-upgrades", log)

    def step_sysctl(log):
        Path("/etc/sysctl.d/99-forward.conf").write_text(sysctl_conf)
        log("Written /etc/sysctl.d/99-forward.conf")
        run("sysctl -p /etc/sysctl.d/99-forward.conf", log)

    def step_wg(log):
        Path("/etc/wireguard/wg0.conf").write_text(wg_conf)
        os.chmod("/etc/wireguard/wg0.conf", 0o600)
        log("Written /etc/wireguard/wg0.conf (mode 600)")

    def step_nft(log):
        Path("/etc/nftables.conf").write_text(nft_conf)
        log("Written /etc/nftables.conf")
        run("systemctl enable nftables", log)
        run("systemctl restart nftables", log)

    def step_wg_start(log):
        run("systemctl enable wg-quick@wg0", log)
        run("systemctl start wg-quick@wg0", log)
        run("wg show", log)

    def step_ssh(log):
        sshd = Path("/etc/ssh/sshd_config")
        text = sshd.read_text()
        def replace_or_append(cfg, key, value):
            pattern = rf"^#?{re.escape(key)}\s+.*"
            line = f"{key} {value}"
            if re.search(pattern, cfg, re.MULTILINE):
                return re.sub(pattern, line, cfg, flags=re.MULTILINE)
            return cfg + f"\n{line}\n"
        text = replace_or_append(text, "Port",            ssh_port)
        text = replace_or_append(text, "PermitRootLogin", "yes")
        text = replace_or_append(text, "PubkeyAuthentication", "yes")
        text = replace_or_append(text, "PasswordAuthentication", "no")
        sshd.write_text(text)
        log(f"sshd_config updated: Port={ssh_port}, PermitRootLogin=yes, PubkeyAuthentication=yes")
        run("systemctl restart ssh || systemctl restart sshd", log, check=False)

    steps = [
        ("apt install wireguard nftables unattended-upgrades", step_apt),
        ("Write sysctl (ip_forward + rp_filter)",              step_sysctl),
        ("Write /etc/wireguard/wg0.conf",                      step_wg),
        ("Write /etc/nftables.conf & restart",                 step_nft),
        ("Enable & start wg-quick@wg0",                        step_wg_start),
        ("Configure sshd (shame-edges profile)",               step_ssh),
    ]

    runner = StepRunner(stdscr, steps, "Installing Edge…")
    ok = runner.run()

    next_steps = [
        "✓ Edge installation complete!" if ok else "✗ Edge installation had errors — check log above.",
        "",
        f"Edge public key (give this to your client):",
        f"  {edge_pub}",
        "",
        "Next steps:",
        "  1. On the CLIENT, run this script and choose 'Client'.",
        f"  2. Give the client operator your edge pub key above.",
        f"  3. Get the client's public key and fill it in:",
        "       /etc/wireguard/wg0.conf  →  [Peer] PublicKey = <client.pub>",
        "     then: systemctl restart wg-quick@wg0",
        "",
        f"  4. SSH is now on port {ssh_port} — update your client if needed!",
        "     (Don't lock yourself out before reconnecting on the new port.)",
        "",
        "Config files written:",
        "  /etc/wireguard/wg0.conf",
        "  /etc/nftables.conf",
        "  /etc/sysctl.d/99-forward.conf",
    ]
    info_screen(stdscr, "Edge Done 🗞️" if ok else "Edge — Errors", next_steps,
                C_SUCCESS if ok else C_ERROR)

# ─── Client installer ─────────────────────────────────────────────────────────

def install_client(stdscr):
    params = collect_client_params(stdscr)

    vps_ip   = params["vps_ip"]
    wg_port  = params["wg_port"]
    edge_pub = params["edge_pub"].strip()

    # Generate keys (pure Python, no wg binary needed)
    client_priv, client_pub = wg_genkey()
    wg_conf = make_wg_client(client_priv, edge_pub or "<EDGE_PUB_KEY>", vps_ip, wg_port)

    summary = [
        f"Role         : Client (Home server)",
        f"VPS IP       : {vps_ip}",
        f"WG port      : {wg_port}",
        f"Edge pub key : {edge_pub or '(fill in manually)'}",
        f"Client pub   : {client_pub}",
        "",
        "── /etc/wireguard/wg0.conf ────────────────────────────",
        *wg_conf.splitlines(),
        "",
        "Note: AllowedIPs = 0.0.0.0/0 routes ALL traffic through the tunnel.",
        "This is intentional for the DSNAT setup.",
    ]

    if not confirm_screen(stdscr, "Client — Review & Confirm", summary):
        return

    def step_apt(log):
        run("apt-get update -qq", log)
        run("apt-get install -y wireguard", log)

    def step_rm_tailscale(log):
        r = run("dpkg -l tailscale 2>/dev/null | grep '^ii'", log, check=False)
        if r.returncode == 0:
            run("apt-get purge -y tailscale", log)
            log("Tailscale removed (hehe)")
        else:
            log("Tailscale not installed, nothing to purge")

    def step_wg(log):
        Path("/etc/wireguard/wg0.conf").write_text(wg_conf)
        os.chmod("/etc/wireguard/wg0.conf", 0o600)
        log("Written /etc/wireguard/wg0.conf (mode 600)")

    def step_wg_start(log):
        run("systemctl enable wg-quick@wg0", log)
        run("systemctl start wg-quick@wg0", log)
        run("wg show", log)

    steps = [
        ("apt install wireguard",         step_apt),
        ("Purge tailscale (if present)",  step_rm_tailscale),
        ("Write /etc/wireguard/wg0.conf", step_wg),
        ("Enable & start wg-quick@wg0",   step_wg_start),
    ]

    runner = StepRunner(stdscr, steps, "Installing Client…")
    ok = runner.run()

    next_steps = [
        "✓ Client installation complete!" if ok else "✗ Client installation had errors.",
        "",
        "Client public key (give this to the EDGE operator):",
        f"  {client_pub}",
        "",
        "The edge operator needs to add this to /etc/wireguard/wg0.conf:",
        "  [Peer]",
        f"  PublicKey = {client_pub}",
        "  AllowedIPs = 10.44.0.2/32",
        "",
        "Then on the edge: systemctl restart wg-quick@wg0",
        "",
        "Verify tunnel once both sides are up:",
        "  ping 10.44.0.1   (from client → should reach edge)",
        "  wg show          (check latest-handshake)",
        "",
        "Your traffic now exits via the VPS IP. 🎉",
    ]
    info_screen(stdscr, "Client Done 🗞️" if ok else "Client — Errors", next_steps,
                C_SUCCESS if ok else C_ERROR)

# ─── Root check ───────────────────────────────────────────────────────────────

def check_root():
    if os.geteuid() != 0:
        print("dsnat-setup.py must be run as root (sudo python3 dsnat-setup.py)")
        sys.exit(1)

# ─── Main ─────────────────────────────────────────────────────────────────────

def main(stdscr):
    init_colors()
    curses.curs_set(0)
    stdscr.keypad(True)

    role = select_role(stdscr)
    if role is None:
        return

    if role == "edge":
        install_edge(stdscr)
    else:
        install_client(stdscr)

if __name__ == "__main__":
    check_root()
    try:
        curses.wrapper(main)
    except KeyboardInterrupt:
        print("\nAborted.")