#!/usr/bin/env python3
"""RapidoDNS for Linux and macOS.

The same idea as the Windows app, in the form those two systems already
expect: a program you run, not a window you keep open.

  1. A resolver on 127.0.0.1:53. Questions leave inside an ordinary HTTPS
     request to our own site, which is the only shape that reliably gets out.

     This was TCP to port 53 at first, on the reasoning that the boxes
     injecting forged answers do it over UDP and cannot beat a question that
     is never asked over UDP. The reasoning was sound and the conclusion was
     wrong: on the networks this exists for, TCP to port 53 does not get out
     either. Every lookup timed out -- including the machine's own
     connectivity check -- so starting it took the whole machine off the
     internet, which is a far worse failure than the one it was fixing.

     TCP is still tried if HTTPS cannot be reached, and if neither can answer
     a single test question, nothing on the machine is changed at all.

  2. For the hostnames of the services you chose, it answers with a stand-in
     address of its own on 127.0.1.x instead of the real one, so the
     connection comes back here.

  3. On each stand-in it listens on the ports those services use, and carries
     what arrives inside one HTTPS connection to dns.rapidoserver.com. The
     name of the site being visited never appears on the wire; the only name
     there is ours, and the inner TLS is still end to end -- this program
     cannot read it and neither can anything between.

Nothing here needs a kernel module, a tun device, or a package. Python 3.7 and
the standard library, which every one of these machines already has.

  sudo ./rapidodns --list
  sudo ./rapidodns discord steam
  sudo ./rapidodns --all

Ctrl-C puts your DNS back. So does the program dying: the resolver settings are
restored from an atexit handler and from the signal handlers, and the file it
saves them in is left behind on disk so a machine that lost power mid-session
can be put right with --restore.
"""

import argparse
import atexit
import errno
import json
import os
import platform
import re
import select
import shutil
import signal
import socket
import ssl
import struct
import subprocess
import sys
import threading
import time

VERSION = "1.29.0"

HERE = os.path.dirname(os.path.abspath(__file__))
STATE = "/var/lib/rapidodns"
SAVED = os.path.join(STATE, "resolv.saved")
DEVICE = os.path.join(STATE, "device.txt")


# --------------------------------------------------------------- the settings

def load_json(name, embedded):
    """A file beside the program wins; otherwise what was shipped with it."""
    path = os.path.join(HERE, name)
    if os.path.exists(path):
        try:
            with open(path, encoding="utf-8-sig") as f:
                return json.load(f)
        except Exception as e:
            warn("could not read %s (%s); using the built-in copy" % (name, e))
    return embedded


# 443 is always here because every redirected hostname resolves to us, so this
# has to carry their ordinary HTTPS as well. The rest are the ports the game
# launchers turned out to need, found by watching rather than by guessing:
# Battle.net answers on 1119, EA's messaging on 8095 and 9000, its desktop
# config partly on 44325.
DEFAULT_PORTS = [443, 1119, 8095, 9000, 42127, 42230, 44325]

# The shape of the settings, not the settings. The tunnel token used to be
# right here as a fallback, which meant this file could not be put under
# version control without putting the token there too -- permanently, because
# git keeps every version of everything.
#
# The token ships in servers.json beside this program, which is what the
# released tarball contains. Running from a checkout without one is a mistake
# worth naming rather than limping past.
DEFAULT_SERVERS = [{
    "Name": "RapidoDNS — Iran",
    "Dns": "176.120.17.28",
    "Relay": "176.120.17.71",
    "Api": "https://dns.rapidoserver.com/api",
    "Tunnel": "dns.rapidoserver.com",
    "TunnelToken": "",
}]


def say(msg):
    # Piping into head closes the far end mid-write; that is the pipe doing
    # what pipes do, not this program failing, and a traceback for it is noise.
    try:
        sys.stdout.write(msg + "\n")
        sys.stdout.flush()
    except BrokenPipeError:
        try:
            sys.stdout.close()
        except Exception:
            pass
        os._exit(0)


def warn(msg):
    sys.stderr.write("  " + msg + "\n")
    sys.stderr.flush()


# --------------------------------------------------------- questions over TLS

class Doh:
    """DNS inside an ordinary HTTPS request to our own site. RFC 8484.

    Reached by address, never by name, because this is the thing that answers
    names -- looking one up to get here would be a circle with nothing at the
    bottom of it. The certificate is still checked against the hostname, so
    the connection is as safe as any other to that site.

    One connection, kept open. A TLS handshake per question would add a
    quarter of a second to every lookup on the machine, which is the sort of
    slowness people blame on the product for ever afterwards.
    """

    def __init__(self, ip, host, log):
        self.ip = ip
        self.host = host
        self.log = log
        self.sock = None
        self.lock = threading.Lock()

    def _open(self):
        raw = socket.create_connection((self.ip, 443), timeout=10)
        raw.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.sock = ssl.create_default_context().wrap_socket(
            raw, server_hostname=self.host)

    def _drop(self):
        try:
            if self.sock is not None:
                self.sock.close()
        except Exception:
            pass
        self.sock = None

    def ask(self, wire, timeout=8):
        with self.lock:
            last = None
            # Twice: a kept-open connection is closed by the far end eventually,
            # and the question that discovers it should not be the one that
            # fails.
            for attempt in (0, 1):
                try:
                    if self.sock is None:
                        self._open()
                    self.sock.settimeout(timeout)

                    head = ("POST /dns-query HTTP/1.1\r\n"
                            "Host: %s\r\n"
                            "Content-Type: application/dns-message\r\n"
                            "Accept: application/dns-message\r\n"
                            "Content-Length: %d\r\n"
                            "Connection: keep-alive\r\n\r\n" % (self.host, len(wire)))
                    self.sock.sendall(head.encode("ascii") + wire)

                    status = read_line(self.sock)
                    if "200" not in status:
                        raise IOError("the resolver said: " + status.strip())

                    length = -1
                    closing = False
                    while True:
                        line = read_line(self.sock).strip()
                        if not line:
                            break
                        low = line.lower()
                        if low.startswith("content-length:"):
                            length = int(line.split(":", 1)[1].strip())
                        elif low.startswith("connection:") and "close" in low:
                            closing = True
                    if length < 0:
                        raise IOError("no length on the answer")
                    body = recv_exactly(self.sock, length)
                    # Honoured rather than ignored: reusing a connection the
                    # far end has said it is closing gets one good answer and
                    # then a failure on the next question.
                    if closing:
                        self._drop()
                    return body
                except Exception as e:
                    last = e
                    self._drop()
            raise last

    def close(self):
        with self.lock:
            self._drop()


# ------------------------------------------------------------- the DNS server

class Question:
    """Just enough of a DNS message to know what was asked and to answer it."""

    def __init__(self, data):
        self.data = data
        self.name = None
        self.qtype = None
        self.ok = False
        try:
            if len(data) < 12:
                return
            qd = struct.unpack(">H", data[4:6])[0]
            if qd < 1:
                return
            i = 12
            parts = []
            hops = 0
            while True:
                if i >= len(data):
                    return
                n = data[i]
                if n == 0:
                    i += 1
                    break
                # A pointer in a question is malformed, but a loop here would
                # be a way to hang this process from outside it.
                if n & 0xC0:
                    hops += 1
                    if hops > 1:
                        return
                    return
                i += 1
                if i + n > len(data):
                    return
                parts.append(data[i:i + n].decode("idna" if False else "latin-1"))
                i += n
            if i + 4 > len(data):
                return
            self.qtype = struct.unpack(">H", data[i:i + 2])[0]
            self.qend = i + 4
            self.name = ".".join(parts).lower()
            self.ok = True
        except Exception:
            self.ok = False

    def answer_a(self, ip, ttl=60):
        """The question echoed back with one A record bolted on."""
        h = bytearray(self.data[:self.qend])
        h[2] = 0x81          # response, recursion desired
        h[3] = 0x80          # recursion available, no error
        struct.pack_into(">H", h, 6, 1)   # one answer
        struct.pack_into(">H", h, 8, 0)   # no authority
        struct.pack_into(">H", h, 10, 0)  # no additional
        rr = b"\xc0\x0c" + struct.pack(">HHIH", 1, 1, ttl, 4)
        rr += socket.inet_aton(ip)
        return bytes(h) + rr


class Resolver:
    """127.0.0.1:53, answering from the stand-ins or from our own resolver."""

    def __init__(self, upstream, engine, log, doh=None):
        self.upstream = upstream
        self.engine = engine            # asked whether to stand in, per answer
        self.log = log
        self.doh = doh
        self._doh_complained = False
        self.sock = None
        self.stop = threading.Event()
        self.asked = 0
        self.stood_in = 0

    def start(self):
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind(("127.0.0.1", 53))
        s.settimeout(0.5)
        self.sock = s
        t = threading.Thread(target=self._serve, daemon=True)
        t.start()

    def _serve(self):
        while not self.stop.is_set():
            try:
                data, addr = self.sock.recvfrom(4096)
            except socket.timeout:
                continue
            except OSError:
                return
            threading.Thread(target=self._one, args=(data, addr), daemon=True).start()

    def _one(self, data, addr):
        q = Question(data)
        self.asked += 1
        try:
            answer = self._upstream(data)

            # Ask first, substitute after.
            #
            # The first version worked the other way round: it pre-allocated one
            # stand-in per configured suffix and answered from that table. Every
            # name under a suffix therefore shared one address, and the relay on
            # it dialled the suffix itself -- so a lookup of
            # api.live.prod.thehelldiversgame.com came back on the same address
            # as gameguard.thehelldiversgame.com, and both were carried to
            # "thehelldiversgame.com", which has no address at all. Discord and
            # Steam hid it, because their bare names do resolve.
            #
            # Now the real answer decides. If our own proxy answered for this
            # name, it is one we carry and it gets an address of its own; if
            # anything else answered, the reply goes back untouched and this
            # program stays out of a path it has no business being on.
            if q.ok and q.qtype == 1:
                real = first_address(answer)
                here = self.engine.standin_for(q.name, real) if real else None
                if here:
                    self.stood_in += 1
                    self.sock.sendto(q.answer_a(here), addr)
                    return

            self.sock.sendto(answer, addr)
        except Exception as e:
            self.log("  lookup failed for %s: %s" % (q.name or "?", e))

    def _upstream(self, data):
        """Over HTTPS where that works, and over TCP where it does not.

        This used to be TCP only, on the reasoning that the boxes injecting
        forged answers do it over UDP and cannot beat a question that is never
        asked over UDP. The reasoning was right and the conclusion was wrong:
        on the network this is built for, TCP to port 53 does not get out
        either. Every lookup timed out, including the machine's own
        connectivity check, so connecting took the whole machine off the
        internet -- which is a far worse failure than the one it was fixing.

        Inside real HTTPS it is ordinary web traffic to an ordinary site and it
        simply arrives. The Windows client found this out months ago and says
        so in its log every time it starts; this one was never told.
        """
        if self.doh is not None:
            try:
                return self.doh.ask(data)
            except Exception as e:
                if not self._doh_complained:
                    self._doh_complained = True
                    self.log("  DNS over HTTPS failed (%s) — falling back to TCP" % e)

        return self._upstream_tcp(data)

    def _upstream_tcp(self, data):
        s = socket.create_connection((self.upstream, 53), timeout=8)
        try:
            s.sendall(struct.pack(">H", len(data)) + data)
            head = recv_exactly(s, 2)
            n = struct.unpack(">H", head)[0]
            return recv_exactly(s, n)
        finally:
            s.close()

    def close(self):
        self.stop.set()
        try:
            self.sock.close()
        except Exception:
            pass


def dns_question(name, qtype=1):
    """A wire-format A question, for probing whether anything answers at all."""
    q = struct.pack(">HHHHHH", 0x7A7A, 0x0100, 1, 0, 0, 0)
    for part in name.split("."):
        q += bytes([len(part)]) + part.encode("ascii")
    return q + b"\x00" + struct.pack(">HH", qtype, 1)


def first_address(msg):
    """The first A record in a reply, or None.

    Walked properly rather than read off the tail: a reply often carries a CNAME
    chain and several records, and the one that matters is not reliably last.
    """
    try:
        answers = struct.unpack(">H", msg[6:8])[0]
        if answers < 1:
            return None
        i = 12
        # Step over the question section.
        for _ in range(struct.unpack(">H", msg[4:6])[0]):
            while i < len(msg) and msg[i]:
                if msg[i] & 0xC0:
                    i += 1
                    break
                i += msg[i] + 1
            i += 5                       # the final 0, then qtype and qclass
        for _ in range(answers):
            if i + 12 > len(msg):
                return None
            if msg[i] & 0xC0:
                i += 2
            else:
                while i < len(msg) and msg[i]:
                    i += msg[i] + 1
                i += 1
            rtype = struct.unpack(">H", msg[i:i + 2])[0]
            rdlen = struct.unpack(">H", msg[i + 8:i + 10])[0]
            i += 10
            if rtype == 1 and rdlen == 4:
                return socket.inet_ntoa(msg[i:i + 4])
            i += rdlen
    except Exception:
        pass
    return None


def recv_exactly(sock, n):
    buf = b""
    while len(buf) < n:
        part = sock.recv(n - len(buf))
        if not part:
            raise IOError("the connection ended early")
        buf += part
    return buf


# ---------------------------------------------------------------- the tunnel

class Tunnel:
    """One HTTPS connection to our own site per connection carried."""

    def __init__(self, host, entrance_ip, token, log):
        self.host = host
        self.ip = entrance_ip
        self.token = token
        self.log = log
        ctx = ssl.create_default_context()
        # The certificate is checked against the name we are asking for; what
        # is deliberately not required is that the address it resolves to came
        # from a resolver we do not control.
        self.ctx = ctx

    def open_twice(self, dst_host, dst_port):
        """One second chance at the entrance.

        Measured from a machine outside the server, about one entrance
        connection in a hundred and twenty is closed part-way through the
        handshake. It is not a pattern -- the same host works either side of it
        -- but with no retry anywhere it became a game that could not reach its
        login server, for no reason visible to anyone.

        Only here, where nothing has been sent yet. Once the application's own
        bytes are in the tunnel there is no second attempt: the request cannot
        be replayed, and retrying then would turn a visible failure into a
        corrupted one.
        """
        try:
            return self.open(dst_host, dst_port)
        except Exception as first:
            time.sleep(0.12)
            try:
                return self.open(dst_host, dst_port)
            except Exception:
                raise first

    def open(self, dst_host, dst_port, timeout=15):
        raw = socket.create_connection((self.ip, 443), timeout=timeout)
        raw.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        try:
            tls = self.ctx.wrap_socket(raw, server_hostname=self.host)
        except Exception:
            raw.close()
            raise
        req = ("GET /t HTTP/1.1\r\n"
               "Host: %s\r\n"
               "Connection: Upgrade\r\n"
               "Upgrade: rapido\r\n"
               "X-Rapido-Token: %s\r\n"
               "X-Rapido-Dst: %s:%d\r\n\r\n" % (self.host, self.token, dst_host, dst_port))
        tls.sendall(req.encode("ascii"))
        status = read_line(tls)
        if "101" not in status:
            tls.close()
            raise IOError("the tunnel refused %s:%d — %s" % (dst_host, dst_port, status.strip()))
        while read_line(tls).strip():
            pass
        return tls


def read_line(sock):
    out = b""
    while not out.endswith(b"\n"):
        c = sock.recv(1)
        if not c:
            break
        out += c
        if len(out) > 4096:
            break
    return out.decode("latin-1")


# ----------------------------------------------------------------- the relay

class Relay:
    """One listening port on one stand-in address."""

    def __init__(self, bind_ip, port, owner, log, dst_port=None):
        self.bind_ip = bind_ip
        self.port = port
        # Normally the far side is the same port we are listening on. --check
        # listens on a port the machine handed out and still has to ask for 443,
        # which the first version did not -- it asked the tunnel for
        # "discord.com:41732" and got a connection that closed at once.
        self.dst_port = dst_port or port
        self.owner = owner            # stand-in ip -> real hostname
        self.log = log
        self.sock = None
        self.stop = threading.Event()
        self.carried = 0
        self.failed = 0

    def start(self):
        """Bind here, on the caller's thread, before returning.

        This used to be a Thread whose run() did the bind, so start() returned
        while the socket was still unbound -- and the resolver, which starts a
        relay and immediately answers the lookup, handed out an address that
        nothing was listening on yet. One connection in a set of ten arrived
        before its own listener existed.
        """
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            s.bind((self.bind_ip, self.port))
        except OSError as e:
            self.log("  cannot listen on %s:%d — %s" % (self.bind_ip, self.port, e))
            return False
        s.listen(64)
        s.settimeout(0.5)
        self.sock = s
        threading.Thread(target=self._accept, daemon=True).start()
        return True

    def _accept(self):
        while not self.stop.is_set():
            try:
                conn, _ = self.sock.accept()
            except socket.timeout:
                continue
            except OSError:
                return
            threading.Thread(target=self._carry, args=(conn,), daemon=True).start()

    def _carry(self, conn):
        host = self.owner.get(self.bind_ip)
        if not host:
            conn.close()
            return
        try:
            far = ENGINE.tunnel.open_twice(host, self.dst_port)
        except Exception as e:
            self.failed += 1
            self.log("  %s:%d did not open — %s" % (host, self.dst_port, e))
            conn.close()
            return
        self.carried += 1
        pump(conn, far)

    def close(self):
        self.stop.set()
        try:
            self.sock.close()
        except Exception:
            pass


def pump(a, b):
    """Copy both ways until both are done.

    Two things here were wrong, and together they produced a failure that
    looked like the network: roughly one connection in three failing its TLS
    handshake, on a relay that had just served two others perfectly.

    The first was shutting BOTH sockets down as soon as EITHER direction
    finished. A client that has sent its request and is waiting for the answer
    has finished sending; that is not a reason to tear down the reply. The
    correct signal is a half-close of the far side, which the other end reads
    as a clean end of stream.

    The second was worse. The sockets were closed as soon as the calling
    thread's copy finished, while the other copier could still be blocked
    inside recv() on one of them. Closing frees the descriptor, the kernel
    hands the very same number to the next accept(), and that stranded thread
    then reads the NEXT connection's bytes -- swallowing a ClientHello that had
    nothing to do with it. Nothing in any log could show that, because from
    every side involved it looks like a connection that simply stopped.

    So: half-close to signal, and nothing is closed until both copiers are done
    with it.
    """
    def one(src, dst):
        try:
            while True:
                data = src.recv(65536)
                if not data:
                    break
                dst.sendall(data)
        except Exception:
            pass
        finally:
            # "I have nothing more to send" -- not "we are finished".
            try:
                dst.shutdown(socket.SHUT_WR)
            except Exception:
                pass

    t = threading.Thread(target=one, args=(a, b), daemon=True)
    t.start()
    one(b, a)
    t.join(timeout=30)
    for s in (a, b):
        try:
            s.close()
        except Exception:
            pass


# ------------------------------------------------------------ the system DNS

class SystemDns:
    """Points the machine here, and puts it back.

    Three arrangements, because these machines have three. The one that is
    chosen is written down along with what it used to be, so restoring does not
    depend on guessing the same way twice.
    """

    def __init__(self, log):
        self.log = log
        self.how = None
        self.before = None

    def take(self):
        os.makedirs(STATE, exist_ok=True)
        if sys.platform == "darwin":
            self._take_macos()
        elif os.path.islink("/etc/resolv.conf") and "systemd" in os.path.realpath("/etc/resolv.conf"):
            self._take_resolvectl()
        else:
            self._take_resolvconf()
        with open(SAVED, "w", encoding="utf-8") as f:
            json.dump({"how": self.how, "before": self.before}, f)

    # -- macOS: every network service, by name
    def _take_macos(self):
        names = self._macos_services()
        self.before = {}
        for n in names:
            out = run(["networksetup", "-getdnsservers", n])
            self.before[n] = [] if "There aren't any" in out else out.split()
            run(["networksetup", "-setdnsservers", n, "127.0.0.1"])
        self.how = "macos"
        self.log("  DNS set on: " + ", ".join(names))

    def _macos_services(self):
        out = run(["networksetup", "-listallnetworkservices"])
        return [l.strip() for l in out.splitlines()[1:]
                if l.strip() and not l.startswith("*")]

    # -- systemd-resolved
    def _take_resolvectl(self):
        link = self._default_link()
        if not link:
            raise RuntimeError("could not tell which interface carries the default route")
        self.before = {"link": link}
        run(["resolvectl", "dns", link, "127.0.0.1"])
        run(["resolvectl", "domain", link, "~."])
        self.how = "resolvectl"
        self.log("  DNS set on %s through systemd-resolved" % link)

    def _default_link(self):
        out = run(["ip", "route", "show", "default"])
        m = re.search(r"\bdev\s+(\S+)", out)
        return m.group(1) if m else None

    # -- a plain file
    def _take_resolvconf(self):
        with open("/etc/resolv.conf", "rb") as f:
            self.before = f.read().decode("utf-8", "replace")
        shutil.copy2("/etc/resolv.conf", os.path.join(STATE, "resolv.conf.orig"))
        with open("/etc/resolv.conf", "w", encoding="utf-8") as f:
            f.write("# put here by rapidodns; the original is in %s\n"
                    "nameserver 127.0.0.1\n" % STATE)
        self.how = "resolv.conf"
        self.log("  /etc/resolv.conf now points at 127.0.0.1")

    def give_back(self):
        if self.how is None:
            return
        try:
            if self.how == "macos":
                for name, servers in (self.before or {}).items():
                    run(["networksetup", "-setdnsservers", name] +
                        (servers if servers else ["empty"]))
            elif self.how == "resolvectl":
                link = (self.before or {}).get("link")
                if link:
                    run(["resolvectl", "revert", link])
            elif self.how == "resolv.conf":
                with open("/etc/resolv.conf", "w", encoding="utf-8") as f:
                    f.write(self.before or "nameserver 1.1.1.1\n")
            self.log("  DNS put back")
        except Exception as e:
            warn("could not put the DNS back automatically: %s" % e)
            warn("run:  sudo %s --restore" % sys.argv[0])
        finally:
            self.how = None
            try:
                os.remove(SAVED)
            except OSError:
                pass


def restore_from_disk():
    """For a machine that was cut off mid-session."""
    if not os.path.exists(SAVED):
        say("nothing to put back — no interrupted session was recorded")
        return 0
    with open(SAVED, encoding="utf-8") as f:
        d = json.load(f)
    s = SystemDns(say)
    s.how = d.get("how")
    s.before = d.get("before")
    s.give_back()
    return 0


def run(cmd):
    try:
        return subprocess.run(cmd, capture_output=True, text=True, timeout=20).stdout
    except Exception as e:
        warn("%s failed: %s" % (cmd[0], e))
        return ""


# ---------------------------------------------------------------- the engine

MAX_STANDINS = 400


class Engine:
    def __init__(self, server, services, log):
        self.server = server
        self.services = services
        self.log = log
        self.tunnel = Tunnel(server["Tunnel"], server["Dns"],
                             server.get("TunnelToken", ""), log)
        self.standins = {}       # the exact hostname asked for -> 127.0.1.x
        self.owner = {}          # 127.0.1.x -> that same exact hostname
        self.relays = []
        self.resolver = None
        self.dns = SystemDns(log)
        self.ports = []
        self._next = 1
        self._lock = threading.Lock()
        # Which addresses mean "our proxy answered for this". Derived from the
        # server's own address rather than written down, so moving the service
        # to another network does not leave a constant here that used to be true.
        self.our_prefix = ".".join(server["Dns"].split(".")[:3]) + "."

    def standin_for(self, name, real_ip):
        """An address of ours for this exact name, or None to leave it alone.

        Only for names our own proxy answered for. Anything resolving elsewhere
        is reached directly -- standing in front of it would put this program on
        a path it has no business being on, and would break it the moment the
        real site moved.
        """
        if not real_ip or not real_ip.startswith(self.our_prefix):
            return None

        with self._lock:
            have = self.standins.get(name)
            if have:
                return have
            if len(self.standins) >= MAX_STANDINS:
                return None

            n = self._next
            self._next += 1
            ip = "127.0.%d.%d" % (1 + n // 250, 1 + n % 250)

            # macOS answers only for 127.0.0.1 unless each address is added.
            if sys.platform == "darwin":
                run(["ifconfig", "lo0", "alias", ip, "up"])

            started = []
            for p in self.ports:
                r = Relay(ip, p, self.owner, self.log)
                r.start()
                started.append(r)
            self.relays.extend(started)

            self.owner[ip] = name
            self.standins[name] = ip
            return ip

    def plan(self):
        """Which ports the relays listen on.

        Not which hostnames -- those are no longer decided in advance. The
        domain lists still matter, but as the answer to "did the customer ask
        for this service", which the server already knows; what reaches here is
        whether our proxy answered, and that is decided one lookup at a time.

        The default ports are always in, whatever was chosen. They used to
        depend on the selection, and the result was that picking Discord
        silently closed the ports EA needs, with nothing on screen connecting
        the two.
        """
        ports = set(DEFAULT_PORTS)
        for s in self.services:
            for p in s.get("ports", []):
                ports.add(int(p))
        return sorted(ports)

    def start(self):
        self.ports = self.plan()
        names = sum(len(s.get("domains", [])) for s in self.services)
        self.log("  %d service(s), %d domain families, ports %s"
                 % (len(self.services), names,
                    ", ".join(str(p) for p in self.ports)))
        self.log("  addresses are handed out as names are looked up")

        self.doh = Doh(self.server["Dns"], self.server["Tunnel"], self.log)
        self.resolver = Resolver(self.server["Dns"], self, self.log, self.doh)

        # Proved before the machine's resolver is touched.
        #
        # Taking over DNS and then discovering that questions cannot be
        # answered is the worst thing this program can do: it does not degrade
        # the connection, it removes it, and the person it happens to is left
        # with a machine that cannot look anything up and no idea why. That is
        # exactly what happened -- every lookup timed out, including Ubuntu's
        # own connectivity check.
        #
        # So one real question is asked first, by whichever route works. If
        # none does, nothing is changed and the reason is said out loud.
        how = self.check_lookups()
        if how is None:
            raise RuntimeError(
                "no way to ask a question: neither HTTPS nor TCP to " +
                self.server["Dns"] + " answered. Nothing on this machine has "
                "been changed.")
        self.log("  questions leave as " + how)

        self.resolver.start()
        self.dns.take()

    def check_lookups(self):
        """Which route can answer a question, if any. Nothing is changed here."""
        probe = dns_question("example.com")

        try:
            answer = self.doh.ask(probe, timeout=8)
            if first_address(answer):
                return "DNS over HTTPS"
        except Exception as e:
            self.log("  DNS over HTTPS did not answer: %s" % e)

        try:
            s = socket.create_connection((self.server["Dns"], 53), timeout=6)
            try:
                s.sendall(struct.pack(">H", len(probe)) + probe)
                n = struct.unpack(">H", recv_exactly(s, 2))[0]
                if first_address(recv_exactly(s, n)):
                    return "TCP to port 53"
            finally:
                s.close()
        except Exception as e:
            self.log("  TCP to port 53 did not answer: %s" % e)

        return None

    def stop(self):
        self.dns.give_back()
        if self.resolver:
            self.resolver.close()
        for r in self.relays:
            r.close()
        if sys.platform == "darwin":
            for ip in self.owner:
                run(["ifconfig", "lo0", "-alias", ip])

    def summary(self):
        carried = sum(r.carried for r in self.relays)
        failed = sum(r.failed for r in self.relays)
        asked = self.resolver.asked if self.resolver else 0
        stood = self.resolver.stood_in if self.resolver else 0
        return ("%d lookups (%d stood in for, over %d addresses), "
                "%d connections carried, %d refused"
                % (asked, stood, len(self.standins), carried, failed))


ENGINE = None


# ------------------------------------------------------------------ the front

def load_services():
    embedded = []
    return load_json("services.json", embedded)


def pick(catalogue, wanted, take_all):
    if take_all:
        return list(catalogue)
    by_id = {s["id"].lower(): s for s in catalogue}
    out, missing = [], []
    for w in wanted:
        s = by_id.get(w.lower())
        (out.append(s) if s else missing.append(w))
    if missing:
        warn("not in the list: " + ", ".join(missing))
        warn("run with --list to see the names")
    return out


def main():
    global ENGINE
    ap = argparse.ArgumentParser(
        prog="rapidodns",
        description="RapidoDNS %s — for Linux and macOS" % VERSION)
    ap.add_argument("services", nargs="*", help="which services to carry")
    ap.add_argument("--all", action="store_true", help="every service in the list")
    ap.add_argument("--list", action="store_true", help="print the list and stop")
    ap.add_argument("--restore", action="store_true",
                    help="put the DNS back after an interrupted session")
    ap.add_argument("--check", action="store_true",
                    help="test the tunnel and stop, without touching the DNS")
    ap.add_argument("--version", action="store_true")
    a = ap.parse_args()

    if a.version:
        say("rapidodns %s on %s" % (VERSION, platform.platform()))
        return 0

    catalogue = load_services()
    if a.list:
        say("%-18s %-26s %-6s %s" % ("id", "name", "kind", "domains"))
        for s in sorted(catalogue, key=lambda s: (s.get("kind", ""), s["id"])):
            say("%-18s %-26s %-6s %d" % (s["id"], s.get("name", ""),
                                         s.get("kind", ""), len(s.get("domains", []))))
        say("")
        say("%d services. Ports %s are always carried; a few services add more."
            % (len(catalogue), ", ".join(str(p) for p in DEFAULT_PORTS)))
        return 0

    if a.restore:
        need_root()
        return restore_from_disk()

    servers = load_json("servers.json", DEFAULT_SERVERS)
    server = servers[0]

    if not server.get("TunnelToken"):
        warn("there is no tunnel token in servers.json, so the tunnel will")
        warn("refuse every connection. The released tarball ships with one;")
        warn("a checkout does not, because it would be in the repository for")
        warn("ever. Copy the token into servers.json beside this program.")
        return 2

    if a.check:
        return check(server)

    chosen = pick(catalogue, a.services, a.all)
    if not chosen:
        ap.print_usage()
        warn("nothing chosen. --all, or names from --list")
        return 2

    need_root()

    say("rapidodns %s" % VERSION)
    say("  server %s  ·  tunnel %s" % (server["Dns"], server["Tunnel"]))
    say("  services: " + ", ".join(s.get("name", s["id"]) for s in chosen))

    ENGINE = Engine(server, chosen, say)
    stopping = threading.Event()

    def bye(*_):
        stopping.set()

    signal.signal(signal.SIGINT, bye)
    signal.signal(signal.SIGTERM, bye)
    atexit.register(lambda: ENGINE.stop())

    try:
        ENGINE.start()
    except Exception as e:
        warn("could not start: %s" % e)
        ENGINE.stop()
        return 1

    say("  ready — leave this running. Ctrl-C puts your DNS back.")
    try:
        while not stopping.is_set():
            stopping.wait(1)
    finally:
        say("")
        say("  " + ENGINE.summary())
        ENGINE.stop()
    return 0


def check(server):
    """Does the tunnel work from here? Nothing on the machine is changed.

    It goes through a real relay on a spare loopback port rather than talking to
    the tunnel directly. The first version wrapped a second TLS session around
    the tunnel's own socket, which Python cannot do -- wrap_socket works on the
    file descriptor, so the inner handshake went out underneath the outer one
    and every site came back with "unexpected message". Nothing was wrong with
    the tunnel; the test was.

    Going through the relay also means this exercises the same path the games
    do, so a pass here is worth something.
    """
    say("rapidodns %s — checking, without changing anything" % VERSION)
    global ENGINE
    ENGINE = Engine(server, [], say)

    bad = 0
    for host, port in (("discord.com", 443), ("store.steampowered.com", 443),
                       ("www.youtube.com", 443)):
        start = time.time()
        # One relay, on a port the machine picked, pointed at this hostname.
        probe = socket.socket()
        probe.bind(("127.0.0.1", 0))
        spare = probe.getsockname()[1]
        probe.close()

        r = Relay("127.0.0.1", spare, {"127.0.0.1": host}, say, dst_port=port)
        if not r.start():
            say("  %-28s could not listen locally" % host)
            bad += 1
            continue
        try:
            raw = socket.create_connection(("127.0.0.1", spare), timeout=25)
            tls = ssl.create_default_context().wrap_socket(raw, server_hostname=host)
            tls.sendall(("HEAD / HTTP/1.1\r\nHost: %s\r\n"
                         "Connection: close\r\nUser-Agent: rapidodns/%s\r\n\r\n"
                         % (host, VERSION)).encode())
            line = read_line(tls).strip()
            ms = int((time.time() - start) * 1000)
            tls.close()
            if line.startswith("HTTP/"):
                say("  %-28s %s   (%d ms)" % (host, line, ms))
            else:
                say("  %-28s opened but said nothing back" % host)
                bad += 1
        except Exception as e:
            say("  %-28s failed — %s" % (host, e))
            bad += 1
        finally:
            r.close()

    say("")
    say("  " + ("everything answered" if bad == 0 else "%d of 3 did not" % bad))
    return 0 if bad == 0 else 1


def need_root():
    if os.geteuid() != 0:
        warn("this has to run as root: it listens on port 53 and changes the "
             "machine's resolver.")
        warn("try:  sudo %s %s" % (sys.argv[0], " ".join(sys.argv[1:])))
        sys.exit(1)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(0)
