#!/usr/bin/env python3
"""Install the shared-key VM connection without installing application dependencies."""
import argparse
import base64
import gzip
import hashlib
import os
from pathlib import Path
import platform
import shlex
import subprocess
import sys
import urllib.request

VERSION = '1.12.0'
CHECKSUMS = {
    'linux_amd64': 'f3f180f1d93aa72cce4e6386f98cc06569a0146fbd65eb4423cf83e6434bcfe6',
    'linux_arm64': '2ec6152cd2c74fe0146d4d79e4e7aa174521368c56e433d55e023a92ea404ec3',
    'darwin_arm64': '707a4b932eea214765146504a0df246cefc415b4297af65a80dd67cf69ba85a9',
}
HOST_KEY = 'ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIBq0RE3paxxbHyuqPYy50k8KkpZT05SrIRgeKVqhpC7T'
KEY_FINGERPRINT = 'SHA256:6OKOhNI33nhZmOZ2dB7gQSw5Y88r+1X+Q7xnVS8gDvE'
TUNNEL_FINGERPRINT = 'J6/chpuE19+5lFL3hKbVkU4fMx8KheAnchyNzBlB3N0='
TUNNEL_URL = 'https://oracle-bloom-dev.tailc74775.ts.net'


def private_file(path, data, mode=0o600):
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode)
    with os.fdopen(fd, 'wb') as f:
        f.write(data)
    path.chmod(mode)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--transport', choices=['direct', 'https'], required=True)
    parser.add_argument('--key-file', type=Path)
    parser.add_argument('--tunnel-url', default=TUNNEL_URL)
    args = parser.parse_args()
    if not args.tunnel_url.startswith('https://') and not args.tunnel_url.startswith('http://127.0.0.1:'):
        parser.error('Tunnel URL must use HTTPS (HTTP is allowed only for a loopback test).')
    os.umask(0o077)
    directory = Path.home() / '.config/broccoli/env' / ('cloud-vm-' + args.transport)
    directory.mkdir(parents=True, exist_ok=True, mode=0o700)
    key_file = directory / 'id_ed25519'
    if args.key_file:
        key = args.key_file.read_bytes()
    elif os.environ.get('CLOUD_VM_SSH_PRIVATE_KEY_B64'):
        key = base64.b64decode(os.environ.pop('CLOUD_VM_SSH_PRIVATE_KEY_B64'), validate=True)
    elif key_file.exists():
        key = key_file.read_bytes()
    else:
        parser.error('Set the CLOUD_VM_SSH_PRIVATE_KEY_B64 runtime secret once in this provider environment.')
    key = key.strip() + b'\n'
    private_file(key_file, key)
    fingerprint = subprocess.check_output(['ssh-keygen', '-lf', str(key_file)], text=True).split()[1]
    if fingerprint != KEY_FINGERPRINT:
        raise SystemExit('The shared cloud key fingerprint does not match. Connection refused.')
    private_file(directory / 'known_hosts', ('cloud-dev-vm ' + HOST_KEY + '\n').encode())
    proxy_line = ''
    if args.transport == 'https':
        machine = {'aarch64': 'arm64', 'arm64': 'arm64', 'x86_64': 'amd64'}.get(platform.machine())
        target = platform.system().lower() + '_' + str(machine)
        if target not in CHECKSUMS:
            raise SystemExit('Unsupported client platform: ' + target)
        binary = directory / 'chisel'
        stamp = directory / 'chisel-version'
        if not binary.exists() or not stamp.exists() or stamp.read_text() != VERSION:
            url = f'https://github.com/jpillora/chisel/releases/download/v{VERSION}/chisel_{VERSION}_{target}.gz'
            with urllib.request.urlopen(url, timeout=60) as response:
                data = response.read()
            if hashlib.sha256(data).hexdigest() != CHECKSUMS[target]:
                raise SystemExit('Chisel checksum mismatch. Nothing executed.')
            private_file(binary, gzip.decompress(data), 0o700)
            private_file(stamp, VERSION.encode())
        proxy = directory / 'proxy.py'
        private_file(proxy, ('''import hashlib, os, subprocess, sys
from pathlib import Path
root = Path(__file__).parent
os.environ['AUTH'] = 'cloud:' + hashlib.sha256((root / 'id_ed25519').read_bytes()).hexdigest()
cmd = [str(root / 'chisel'), 'client', '--fingerprint', ''' + repr(TUNNEL_FINGERPRINT) + ''', '--max-retry-count', '0']
proxy = os.environ.get('HTTPS_PROXY') or os.environ.get('https_proxy')
if proxy:
    cmd += ['--proxy', proxy]
cmd += [''' + repr(args.tunnel_url) + ''', 'stdio:127.0.0.1:22']
# Chisel logs the proxy URL, which can contain credentials.
result = subprocess.run(cmd, stderr=subprocess.DEVNULL)
if result.returncode:
    print('HTTPS SSH transport failed; check endpoint, proxy policy, and tunnel service.', file=sys.stderr)
sys.exit(result.returncode)
''').encode())
        proxy_line = f'    ProxyCommand {shlex.quote(sys.executable)} {shlex.quote(str(proxy))}\n'
    config = directory / 'ssh_config'
    private_file(config, (f'''Host dev-vm
    HostName 68.233.104.110
    User devagent
    IdentityFile "{key_file}"
    IdentitiesOnly yes
    BatchMode yes
    StrictHostKeyChecking yes
    HostKeyAlias cloud-dev-vm
    UserKnownHostsFile "{directory / 'known_hosts'}"
    HostKeyAlgorithms ssh-ed25519
    ConnectTimeout 20
    ServerAliveInterval 30
    ServerAliveCountMax 3
    ControlMaster auto
    ControlPersist 600
    ControlPath "{directory / 'connection'}"
''' + proxy_line).encode())
    bindir = Path.home() / '.local/bin'
    bindir.mkdir(parents=True, exist_ok=True)
    private_file(bindir / 'vm', f'#!/bin/sh\nexec ssh -F {shlex.quote(str(config))} dev-vm "$@"\n'.encode(), 0o700)
    private_file(bindir / 'vm-copy', f'#!/bin/sh\nexec scp -F {shlex.quote(str(config))} "$@"\n'.encode(), 0o700)
    check = 'hostname; id -un; for p in 4000/health 4111/foreman/health 8000/heartbeat; do curl -fsS --max-time 10 -o /dev/null -w "HTTP %{http_code}\\n" http://127.0.0.1:$p || exit; done'
    subprocess.run([str(bindir / 'vm'), check], check=True)
    print('Ready. Use ~/.local/bin/vm "COMMAND" or ~/.local/bin/vm-copy FILE dev-vm:PATH')


if __name__ == '__main__':
    main()