#!/usr/bin/env python3
"""Join Work once. Parallel mode never touches the system Tailscale client."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import re
import shlex
import shutil
import socket
import socketserver
import struct
import subprocess
import sys
import tarfile
import threading
import time
import urllib.request
from urllib.parse import urlsplit

VERSION = '1.104.4'


def daemon_worker(root, port):
    """Keep only this isolated daemon alive; serialize concurrent setup runs."""
    import fcntl
    import signal
    guard = open(root / 'daemon.lock', 'a')
    try:
        fcntl.flock(guard, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        return
    _, tsd = binaries(root)
    child = None
    def stop(*_):
        if child and child.poll() is None:
            child.terminate()
        raise SystemExit(0)
    signal.signal(signal.SIGTERM, stop)
    while True:
        (root / 'ts.sock').unlink(missing_ok=True)
        log = open(root / 'tailscaled.log', 'a')
        child = subprocess.Popen([tsd, '--tun=userspace-networking', '--port=0',
            '--statedir=' + str(root / 'var'), '--state=' + str(root / 'tailscaled.state'),
            '--socket=' + str(root / 'ts.sock'), '--socks5-server=127.0.0.1:' + str(port + 1)],
            stdout=log, stderr=log, env={**os.environ, 'TS_NO_LOGS_NO_SUPPORT': 'true'})
        (root / 'tailscaled.pid').write_text(str(child.pid))
        child.wait()
        log.close()
        time.sleep(2)


def receive(sock, length):
    data = b''
    while len(data) < length:
        part = sock.recv(length - len(data))
        if not part:
            raise ConnectionError('Closed SOCKS connection')
        data += part
    return data


def owned_peer(connection):
    """A loopback listener is not a user boundary on a shared workstation."""
    if connection.family == socket.AF_UNIX and hasattr(socket, 'SO_PEERCRED'):
        _, uid, _ = struct.unpack('3i', connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
        return uid == os.getuid()
    peer, local = connection.getpeername(), connection.getsockname()
    if platform.system() == 'Linux':
        left = '0100007F:%04X' % peer[1]
        right = '0100007F:%04X' % local[1]
        for line in Path('/proc/net/tcp').read_text().splitlines()[1:]:
            fields = line.split()
            if fields[1] == left and fields[2] == right:
                return int(fields[7]) == os.getuid()
        return False
    if platform.system() == 'Darwin':
        result = subprocess.run(['/usr/sbin/lsof', '-nP', '-iTCP:' + str(peer[1]),
            '-sTCP:ESTABLISHED', '-Fun'], capture_output=True, text=True, timeout=5)
        uid = None
        expected = f'127.0.0.1:{peer[1]}->127.0.0.1:{local[1]}'
        for line in result.stdout.splitlines():
            if line.startswith('u'): uid = int(line[1:])
            if line == 'n' + expected: return uid == os.getuid()
    return False


def splice(left, right):
    def pump(source, sink):
        try:
            while chunk := source.recv(65536): sink.sendall(chunk)
        except OSError: pass
        try: sink.shutdown(socket.SHUT_WR)
        except OSError: pass
    sender = threading.Thread(target=pump, args=(left, right), daemon=True)
    sender.start(); pump(right, left); sender.join()


def own_proxy_process(root):
    try:
        pid = int((root / 'proxy.pid').read_text())
        os.kill(pid, 0)
        if platform.system() == 'Linux':
            arguments = Path(f'/proc/{pid}/cmdline').read_bytes().split(b'\0')
            return b'--proxy-worker' in arguments and str(root).encode() in arguments
        command = subprocess.check_output(['ps', '-p', str(pid), '-o', 'command='], text=True)
        return '--proxy-worker' in command and str(root) in command
    except (OSError, ValueError, subprocess.CalledProcessError):
        return False


def relay_worker(root, port, target):
    class Relay(socketserver.BaseRequestHandler):
        def handle(self):
            if not owned_peer(self.request): return
            try:
                with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as upstream:
                    upstream.connect(target)
                    splice(self.request, upstream)
            except OSError: pass
    class Server(socketserver.ThreadingTCPServer):
        allow_reuse_address = True
        daemon_threads = True
    with Server(('127.0.0.1', port), Relay) as server:
        server.serve_forever()


def proxy_worker(root, port):
    """Resolve private names from our netmap, never the system's other tailnet."""
    ts, _ = binaries(root)
    lock, cache = threading.Lock(), {'at': 0, 'hosts': {}, 'private_suffixes': {'mesh'}}
    def resolve(host):
        with lock:
            if time.monotonic() - cache['at'] > 3:
                result = subprocess.run([ts, '--socket=' + str(root / 'ts.sock'), 'debug', 'netmap'],
                                        capture_output=True, text=True, timeout=10)
                if result.returncode:
                    raise ConnectionError('Private network is disconnected')
                network = json.loads(result.stdout)
                cache['hosts'] = {r['Name'].rstrip('.').lower(): r['Value']
                                  for r in network.get('DNS', {}).get('ExtraRecords', []) if r['Type'] == 'A'}
                cache['private_suffixes'].update(name.rsplit('.', 1)[1] for name in cache['hosts'])
                cache['at'] = time.monotonic()
            name = host.rstrip('.').lower()
            if name not in cache['hosts'] and name.rsplit('.', 1)[-1] in cache['private_suffixes']:
                raise ConnectionError('No such private service')
            return cache['hosts'].get(name, host)
    class Proxy(socketserver.BaseRequestHandler):
        def handle(self):
            client = self.request
            upstream = None
            try:
                if not owned_peer(client): return
                client.settimeout(20)
                version, count = receive(client, 2)
                methods = receive(client, count)
                if version != 5 or 0 not in methods:
                    client.sendall(b'\x05\xff'); return
                client.sendall(b'\x05\x00')
                version, command, reserved, family = receive(client, 4)
                if version != 5 or command != 1 or reserved:
                    raise ValueError('Only SOCKS5 CONNECT is supported')
                if family == 3:
                    host = receive(client, receive(client, 1)[0]).decode('ascii')
                    host = resolve(host)
                elif family in (1, 4):
                    host = socket.inet_ntop(socket.AF_INET if family == 1 else socket.AF_INET6, receive(client, 4 if family == 1 else 16))
                else:
                    raise ValueError('Unknown SOCKS address family')
                destination_port = receive(client, 2)
                try:
                    address = b'\x01' + socket.inet_pton(socket.AF_INET, host)
                except OSError:
                    try:
                        address = b'\x04' + socket.inet_pton(socket.AF_INET6, host)
                    except OSError:
                        encoded = host.encode('ascii')
                        address = b'\x03' + bytes([len(encoded)]) + encoded
                upstream = socket.create_connection(('127.0.0.1', port + 1), timeout=20)
                upstream.sendall(b'\x05\x01\x00')
                if receive(upstream, 2) != b'\x05\x00':
                    raise ConnectionError('Isolated SOCKS client refused authentication')
                upstream.sendall(b'\x05\x01\x00' + address + destination_port)
                response = receive(upstream, 4)
                bound = receive(upstream, 4 if response[3] == 1 else 16) if response[3] in (1, 4) else receive(upstream, receive(upstream, 1)[0])
                bound_port = receive(upstream, 2)
                client.sendall(response + bound + bound_port)
                if response[1]: return
                client.settimeout(None); upstream.settimeout(None)
                def pump(source, sink):
                    try:
                        while part := source.recv(65536): sink.sendall(part)
                    except OSError: pass
                    try: sink.shutdown(socket.SHUT_WR)
                    except OSError: pass
                sender = threading.Thread(target=pump, args=(client, upstream), daemon=True)
                sender.start(); pump(upstream, client); sender.join()
            except (OSError, ValueError, KeyError, subprocess.TimeoutExpired):
                try: client.sendall(b'\x05\x04\x00\x01\x00\x00\x00\x00\x00\x00')
                except OSError: pass
            finally:
                if upstream: upstream.close()
    unix = os.environ.get('MESH_PROXY_UNIX')
    base = socketserver.ThreadingUnixStreamServer if unix else socketserver.ThreadingTCPServer
    class Server(base):
        allow_reuse_address = True
        daemon_threads = True
    if unix: Path(unix).unlink(missing_ok=True)
    address = unix if unix else ('127.0.0.1', port)
    with Server(address, Proxy) as server:
        server.serve_forever()


def download(url, target):
    with urllib.request.urlopen(url, timeout=60) as response:
        target.write_bytes(response.read())
    return target


def binaries(root):
    if platform.system() != 'Linux':
        ts, tsd = shutil.which('tailscale'), shutil.which('tailscaled')
        if platform.system() == 'Darwin' and (not ts or not tsd) and shutil.which('brew'):
            brew = shutil.which('brew')
            result = subprocess.run([brew, '--prefix', 'tailscale'], capture_output=True, text=True)
            prefix = result.stdout.strip() if result.returncode == 0 else ''
            folder = Path(prefix) / 'bin'
            if not prefix or not all((folder / name).exists() for name in ('tailscale', 'tailscaled')):
                print('Installing standalone client tools with Homebrew; the existing app/service is not restarted.', flush=True)
                subprocess.run([brew, 'install', '--formula', 'tailscale'], check=True)
                folder = Path(subprocess.check_output([brew, '--prefix', 'tailscale'], text=True).strip()) / 'bin'
            ts, tsd = str(folder / 'tailscale'), str(folder / 'tailscaled')
        if not ts or not tsd:
            raise SystemExit('Parallel mode needs the standalone tailscale and tailscaled binaries. On macOS install the standalone client, or use the app with the custom server URL.')
        return ts, tsd
    architecture = {'x86_64': 'amd64', 'aarch64': 'arm64'}.get(platform.machine())
    if not architecture:
        raise SystemExit('Use the official app on this CPU architecture')
    folder = root / 'bin'
    folder.mkdir(exist_ok=True)
    if not (folder / 'tailscaled').exists():
        url = f'https://pkgs.tailscale.com/stable/tailscale_{VERSION}_{architecture}.tgz'
        archive = download(url, root / 'tailscale.tgz')
        with urllib.request.urlopen(url + '.sha256', timeout=30) as response:
            expected = response.read().decode().split()[0]
        if hashlib.sha256(archive.read_bytes()).hexdigest() != expected:
            raise SystemExit('Downloaded client checksum differs from the official checksum')
        with tarfile.open(archive) as source:
            for name in ('tailscale', 'tailscaled'):
                member = next(m for m in source.getmembers() if m.name.endswith('/' + name))
                target = folder / name
                target.write_bytes(source.extractfile(member).read())
                target.chmod(0o700)
        archive.unlink()
    return str(folder / 'tailscale'), str(folder / 'tailscaled')


def trust(ca):
    if platform.system() == 'Darwin':
        subprocess.run(['security', 'add-trusted-cert', '-r', 'trustRoot', '-k',
                        str(Path.home() / 'Library/Keychains/login.keychain-db'), str(ca)], check=True)
    elif platform.system() == 'Linux':
        if shutil.which('update-ca-certificates'):
            subprocess.run(['sudo', 'install', '-m', '644', str(ca), '/usr/local/share/ca-certificates/work-mesh-prototype.crt'], check=True)
            subprocess.run(['sudo', 'update-ca-certificates'], check=True)
        else:
            print('System trust requires your distribution’s certificate configuration. The CA is at', ca)
        certutil = shutil.which('certutil')
        if certutil:
            for folder in [Path.home() / '.pki/nssdb', Path.home() / '.local/share/pki/nssdb']:
                folder.mkdir(parents=True, exist_ok=True, mode=0o700)
                database = 'sql:' + str(folder)
                if not (folder / 'cert9.db').exists():
                    subprocess.run([certutil, '-N', '-d', database, '--empty-password'], check=True)
                subprocess.run([certutil, '-A', '-d', database, '-n', 'Work Mesh Prototype', '-t', 'C,,', '-i', str(ca)], check=True)
    else:
        raise SystemExit('Install the downloaded CA using your operating system’s certificate manager')


def browser_launcher(root, port, pin, server, join):
    candidates = [shutil.which(name) for name in ['google-chrome', 'google-chrome-stable', 'chromium',
                    'chromium-browser', 'brave-browser', 'microsoft-edge']]
    candidates += ['/Applications/Google Chrome.app/Contents/MacOS/Google Chrome',
                   '/Applications/Chromium.app/Contents/MacOS/Chromium',
                   '/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge',
                   '/Applications/Brave Browser.app/Contents/MacOS/Brave Browser']
    browser = next((path for path in candidates if path and Path(path).is_file()), None)
    if not browser:
        print('For a parallel browser trial, install Chrome, Chromium, Edge or Brave and rerun with --open.')
        return
    command = [browser, '--user-data-dir=' + str(root / 'browser'),
               '--proxy-server=socks5://127.0.0.1:' + str(port),
               '--ignore-certificate-errors-spki-list=' + pin,
               '--no-first-run', '--no-default-browser-check', 'https://work.mesh/']
    # Pinning is limited to this private CA and this separate browser process.
    # No global certificate-error bypass or system certificate change is used.
    launcher = root / 'open-work.py'
    bootstrap = [str(root / 'setup-client.py'), '--server', server, '--join-url', join,
                 '--state', str(root), '--parallel', '--socks-port', str(port)]
    launcher.write_text('#!/usr/bin/env python3\nimport subprocess, sys\n'
        + 'if subprocess.run([sys.executable, *' + repr(bootstrap) + ']).returncode == 0:\n'
        + '    subprocess.Popen(' + repr(command) + ')\n')
    launcher.chmod(0o700)
    if platform.system() == 'Darwin':
        shortcut = root / 'Open Work.command'
        shortcut.write_text('#!/bin/sh\nexec /usr/bin/env python3 ' + shlex.quote(str(launcher)) + '\n')
        shortcut.chmod(0o700)
        print('Mac launcher:', shortcut)
    subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
    print('Opened a separate Work browser. Reopen it with:', launcher)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--server', default='https://control.2.31.109.233.sslip.io', help='Public control server URL')
    parser.add_argument('--join-url', help='Public onboarding URL; inferred from control.* by default')
    parser.add_argument('--state', default=str(Path.home() / '.local/state/work-mesh-client'))
    parser.add_argument('--parallel', action='store_true', help='Run an isolated userspace client with a SOCKS proxy')
    parser.add_argument('--switch', action='store_true', help='Explicitly switch the system Tailscale client to Work')
    parser.add_argument('--trust', action='store_true', help='Install the private CA into OS/browser trust stores')
    parser.add_argument('--open', action='store_true', help='Open a separate Chrome/Chromium/Edge/Brave profile for the parallel trial')
    parser.add_argument('--ca-sha256', help='Optional out-of-band SHA256 fingerprint of the CA certificate DER')
    parser.add_argument('--auth-key-file', help='Optional one-use enrollment key file; browser SSO is the default')
    parser.add_argument('--socks-port', type=int, default=1088)
    parser.add_argument('--stop', action='store_true', help='Stop this isolated client only')
    parser.add_argument('--proxy-worker', action='store_true', help=argparse.SUPPRESS)
    parser.add_argument('--daemon-worker', action='store_true', help=argparse.SUPPRESS)
    parser.add_argument('--relay-worker', action='store_true', help=argparse.SUPPRESS)
    parser.add_argument('--relay-socket', help=argparse.SUPPRESS)
    args = parser.parse_args()
    os.umask(0o077)
    root = Path(args.state).expanduser().resolve()
    root.mkdir(parents=True, exist_ok=True, mode=0o700)
    root.chmod(0o700)
    if args.relay_worker:
        relay_worker(root, args.socks_port, args.relay_socket)
        return
    if args.proxy_worker:
        proxy_worker(root, args.socks_port)
        return
    if args.daemon_worker:
        daemon_worker(root, args.socks_port)
        return
    managed = root / 'setup-client.py'
    if Path(__file__).resolve() != managed:
        managed.write_bytes(Path(__file__).read_bytes())
        managed.chmod(0o700)
    address = urlsplit(args.server)
    if address.scheme != 'https' or not address.hostname or address.path not in ('', '/'):
        raise SystemExit('Use an HTTPS server origin, without a path')
    if args.stop:
        pid = root / 'tailscaled.pid'
        if pid.exists():
            # Use the private socket for shutdown; do not signal a potentially reused PID.
            ts, _ = binaries(root)
            subprocess.run([ts, '--socket=' + str(root / 'ts.sock'), 'down'], check=True)
            print('The parallel client is disconnected; the system client is unchanged.')
        return
    join = args.join_url or args.server.replace('://control.', '://join.', 1)
    ca = download(join.rstrip('/') + '/ca.crt', root / 'ca.crt')
    if b'PRIVATE KEY' in ca.read_bytes():
        raise SystemExit('The server returned a private key instead of a certificate')
    import ssl
    fingerprint = hashlib.sha256(ssl.PEM_cert_to_DER_cert(ca.read_text())).hexdigest()
    if args.ca_sha256 and fingerprint != args.ca_sha256.lower().replace(':', ''):
        raise SystemExit('The CA fingerprint does not match; no trust or enrollment changes were made')
    print('Work CA SHA256:', fingerprint)
    manifest = None
    if args.open:
        with urllib.request.urlopen(join.rstrip('/') + '/client.json', timeout=15) as response:
            manifest = json.load(response)
        if manifest['ca_sha256'] != fingerprint:
            raise SystemExit('The onboarding manifest and certificate do not agree')
    if args.trust:
        trust(ca)
    if args.parallel == args.switch:
        raise SystemExit('Choose --parallel for a trial or --switch for an explicit system-client cutover')
    if args.parallel:
        ts, tsd = binaries(root)
        socket_path = str(root / 'ts.sock')
        if Path(socket_path).exists():
            try:
                with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
                    probe.settimeout(1); probe.connect(socket_path)
            except OSError:
                Path(socket_path).unlink(missing_ok=True)
        if not Path(socket_path).exists():
            log = open(root / 'supervisor.log', 'a')
            process = subprocess.Popen([sys.executable, str(Path(__file__).resolve()), '--server', args.server,
                '--state', str(root), '--socks-port', str(args.socks_port), '--daemon-worker'],
                stdout=log, stderr=log, start_new_session=True)
            (root / 'supervisor.pid').write_text(str(process.pid))
            for _ in range(100):
                if Path(socket_path).exists(): break
                if process.poll() not in (None, 0): raise SystemExit('Isolated client failed; inspect ' + str(root / 'tailscaled.log'))
                time.sleep(.1)
        prefix = [ts, '--socket=' + socket_path]
        flags = ['--accept-dns=false', '--accept-routes=false']
    else:
        ts = shutil.which('tailscale')
        if not ts: raise SystemExit('Install the official Tailscale client first')
        prefix, flags = [ts], []
    command = [*prefix, 'up', '--login-server=' + args.server, *flags]
    if args.auth_key_file:
        command.append('--auth-key=file:' + str(Path(args.auth_key_file).resolve()))
    print('Follow the sign-in link from the client to enroll your own device.', flush=True)
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)
    opened = False
    for line in process.stdout:
        print(line, end='', flush=True)
        match = re.search(r'https://[^\s]+', line)
        if match and not opened and not args.auth_key_file and match[0].startswith(args.server.rstrip('/') + '/'):
            import webbrowser
            webbrowser.open(match[0])
            opened = True
    if process.wait():
        raise SystemExit('Enrollment did not complete; rerun the setup to resume')
    if args.parallel:
        # Userspace clients with accept-dns=false must not resolve through the
        # host's currently active tailnet. This loopback adapter uses ExtraRecords.
        try:
            if os.environ.get('MESH_PROXY_UNIX'):
                with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
                    probe.connect(os.environ['MESH_PROXY_UNIX'])
            else:
                with socket.create_connection(('127.0.0.1', args.socks_port), timeout=1): pass
                if not own_proxy_process(root):
                    raise SystemExit('The chosen SOCKS port belongs to another process. Choose a free --socks-port.')
        except OSError:
            log = open(root / 'proxy.log', 'a')
            worker = subprocess.Popen([sys.executable, str(Path(__file__).resolve()), '--server', args.server,
                '--state', str(root), '--socks-port', str(args.socks_port), '--proxy-worker'],
                stdout=log, stderr=log, start_new_session=True)
            (root / 'proxy.pid').write_text(str(worker.pid))
            for _ in range(50):
                try:
                    if os.environ.get('MESH_PROXY_UNIX'):
                        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
                            probe.connect(os.environ['MESH_PROXY_UNIX'])
                    else:
                        with socket.create_connection(('127.0.0.1', args.socks_port), timeout=.2): pass
                    break
                except OSError: time.sleep(.1)
        command = ['curl', '--socks5-hostname', f'127.0.0.1:{args.socks_port}', '--cacert', str(ca), 'https://work.mesh/']
        print('\nParallel connection ready. Your system Tailscale connection is unchanged.')
        print('Test it:', shlex.join(command))
        print(f'For a trial browser, use SOCKS5 127.0.0.1:{args.socks_port} with proxy DNS enabled.')
        print('This proxy setting belongs to the trial browser; it does not change system DNS.')
        if args.open:
            browser_launcher(root, args.socks_port, manifest['ca_spki_sha256'], args.server, join)
    else:
        print('Connected. Open https://work.mesh/ after installing the downloaded CA.')
        if args.open:
            import webbrowser
            webbrowser.open('https://work.mesh/')


if __name__ == '__main__':
    main()