#!/usr/bin/env python3
"""
PocketMage-3DS iPhone Hotspot <-> Tailscale MagicDNS Relay + FTPD Updater (v2.3)
Run on iPhone (a-Shell / Pythonista / iSH):
  curl -s http://nerv.tail53d7c2.ts.net:60002/relay.py | python3
  (or once cached locally in a-Shell: python3 relay.py)
"""
import ftplib, io, os, socket, struct, sys, threading, time, urllib.request

NERV_CANDIDATES = [
    "nerv.tail53d7c2.ts.net",
    "nerv",
    "100.92.246.59",
    "192.168.1.112",
]
if len(sys.argv) > 1 and sys.argv[1] not in NERV_CANDIDATES:
    NERV_CANDIDATES.insert(0, sys.argv[1])

DS_IP = sys.argv[2] if len(sys.argv) > 2 else "172.20.10.3"
ACTIVE_NERV = NERV_CANDIDATES[0]
BOOT_DATA = b""

def connect_nerv(port, timeout=3.0):
    global ACTIVE_NERV
    order = [ACTIVE_NERV] + [h for h in NERV_CANDIDATES if h != ACTIVE_NERV]
    last_err = None
    for host in order:
        try:
            s = socket.create_connection((host, port), timeout=timeout)
            if host != ACTIVE_NERV:
                ACTIVE_NERV = host
                print(f"[+] Connected to nerv via {ACTIVE_NERV}:{port}")
            return s
        except Exception as e:
            last_err = e
    raise last_err

def fetch_boot_3dsx():
    global BOOT_DATA, ACTIVE_NERV
    for host in [ACTIVE_NERV] + [h for h in NERV_CANDIDATES if h != ACTIVE_NERV]:
        try:
            url = f"http://{host}:60002/boot.3dsx"
            data = urllib.request.urlopen(url, timeout=8).read()
            if len(data) > 32768 and data[:4] == b"3DSX":
                ACTIVE_NERV = host
                BOOT_DATA = data
                try:
                    with open("boot.3dsx", "wb") as f:
                        f.write(data)
                except Exception:
                    pass
                try:
                    script_bytes = urllib.request.urlopen(f"http://{host}:60002/relay.py", timeout=5).read()
                    if len(script_bytes) > 500:
                        with open("relay.py", "wb") as f:
                            f.write(script_bytes)
                except Exception:
                    pass
                print(f"[+] Downloaded boot.3dsx ({len(data)//1024} KB) from {host} & cached locally!")
                return data
        except Exception:
            pass
    if os.path.exists("boot.3dsx"):
        try:
            with open("boot.3dsx", "rb") as f:
                data = f.read()
            if len(data) > 32768:
                BOOT_DATA = data
                print(f"[+] Loaded cached boot.3dsx ({len(data)//1024} KB) from local storage!")
                return data
        except Exception:
            pass
    return b""

def push_ftpd_data(ds_ip, ftp_port, data):
    if not data:
        return False
    try:
        ftp = ftplib.FTP()
        ftp.connect(ds_ip, ftp_port, timeout=6)
        try:
            ftp.login()
        except Exception:
            pass
        print(f"[*] Connected to 3DS ftpd ({ds_ip}:{ftp_port})! Uploading v2.3 ({len(data)//1024} KB)...")
        ftp.storbinary("STOR /boot.3dsx", io.BytesIO(data))
        print("[+] Uploaded /boot.3dsx!")

        try:
            ftp.storbinary("STOR /3ds/pocketmage-3ds/pocketmage-3ds.3dsx", io.BytesIO(data))
            print("[+] Uploaded /3ds/pocketmage-3ds/pocketmage-3ds.3dsx!")
        except Exception:
            pass

        try:
            ftp.delete("/3ds/pocketmage-3ds.3dsx")
            print("[+] Removed old /3ds/pocketmage-3ds.3dsx!")
        except Exception:
            try:
                ftp.storbinary("STOR /3ds/pocketmage-3ds.3dsx", io.BytesIO(data))
            except Exception:
                pass

        hosts_ini = (
            b"iphone-relay|172.20.10.1|ctk||2222|60001\n"
            b"nerv-mdns|nerv.tail53d7c2.ts.net|ctk||22|60001\n"
            b"nerv-tail|100.92.246.59|ctk||22|60001\n"
            b"nerv-lan|192.168.1.112|ctk||22|60001\n"
        )
        try:
            ftp.storbinary("STOR /pocketmage/hosts.ini", io.BytesIO(hosts_ini))
            print("[+] Updated /pocketmage/hosts.ini (iphone-relay + MagicDNS nerv.tail53d7c2.ts.net)!")
        except Exception:
            pass

        try:
            ftp.quit()
        except Exception:
            ftp.close()
        print("==========================================================")
        print(" [SUCCESS] POCKETMAGE-3DS v2.3 (3D + TOUCH) INSTALLED!    ")
        print(" Reboot 3DS or exit ftpd to launch v2.3!                  ")
        print("==========================================================")
        return True
    except Exception:
        return False

def tcp_pipe(src, dst):
    try:
        while True:
            buf = src.recv(16384)
            if not buf: break
            dst.sendall(buf)
    except Exception:
        pass
    finally:
        try: src.close()
        except Exception: pass
        try: dst.close()
        except Exception: pass

def start_tcp_forward(listen_port, target_port):
    printed = False
    while True:
        srv = None
        try:
            srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            srv.bind(("0.0.0.0", listen_port))
            srv.listen(8)
            if not printed:
                print(f"[+] TCP Relay active: 172.20.10.1:{listen_port} -> nerv:{target_port} (MagicDNS auto-failover)")
                printed = True
            while True:
                c, addr = srv.accept()
                try:
                    u = connect_nerv(target_port, timeout=5.0)
                    threading.Thread(target=tcp_pipe, args=(c, u), daemon=True).start()
                    threading.Thread(target=tcp_pipe, args=(u, c), daemon=True).start()
                except Exception:
                    # Fallback: if port 60002 OTA requested while Tailscale is interrupted, serve cached BOOT_DATA!
                    if listen_port == 60002 and BOOT_DATA:
                        try:
                            req = c.recv(1024)
                            if req.startswith(b"PM3OTA"):
                                c.sendall(struct.pack("<4sI", b"PM3O", len(BOOT_DATA)) + BOOT_DATA)
                        except Exception:
                            pass
                    try: c.close()
                    except Exception: pass
        except Exception:
            if srv is not None:
                try: srv.close()
                except Exception: pass
            time.sleep(1.0)

def start_udp_forward(listen_port):
    printed = False
    while True:
        usock = None
        try:
            usock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            usock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            usock.bind(("0.0.0.0", listen_port))
            if not printed:
                print(f"[+] Mosh Relay active: 3DS UDP 172.20.10.1:{listen_port} <-> nerv:60002 (MagicDNS auto-failover)")
                printed = True
            while True:
                pkt, addr = usock.recvfrom(65535)
                if not pkt.startswith(b"PM3D"):
                    continue
                try:
                    with connect_nerv(60002, timeout=2.5) as ts:
                        ts.sendall(pkt)
                        resp = b""
                        while True:
                            chunk = ts.recv(4096)
                            if not chunk:
                                break
                            resp += chunk
                        if resp:
                            usock.sendto(resp, addr)
                except Exception:
                    pass
        except Exception:
            if usock is not None:
                try: usock.close()
                except Exception: pass
            time.sleep(1.0)

if __name__ == "__main__":
    print("[*] Resolving nerv via MagicDNS (nerv.tail53d7c2.ts.net / nerv / 100.92.246.59)...")
    fetch_boot_3dsx()

    threading.Thread(target=start_tcp_forward, args=(60002, 60002), daemon=True).start()
    threading.Thread(target=start_tcp_forward, args=(2222,  22),    daemon=True).start()
    threading.Thread(target=start_udp_forward, args=(60001,),       daemon=True).start()

    pushed = False
    print(f"[*] Listening for 3DS ftpd on {DS_IP}:5000 AND relaying Mosh/OTA...")
    try:
        while True:
            if not BOOT_DATA:
                fetch_boot_3dsx()
            if BOOT_DATA and not pushed:
                if push_ftpd_data(DS_IP, 5000, BOOT_DATA):
                    pushed = True
            time.sleep(2.5)
    except KeyboardInterrupt:
        pass
