#!/usr/bin/env python3
"""Install/update a local release binary on Linux/systemd. No builds.

First install downloads the PDDL user-country GeoIP CSV files from the
sapics/ip-location-db GitHub releases and verifies publisher SHA-256 files.
Updates compare the installed GeoIP files with current publisher checksums and
stage only stale assets before service downtime. Updates stop the service before
copying all private data, refresh the managed systemd unit, switch an immutable
release symlink, check health/version and restore code, data and unit on failure.
Use --plan to inspect operations without changing the host.
"""
from __future__ import annotations
import argparse
import fcntl
import hashlib
import http.client
import ipaddress
import json
import os
from pathlib import Path
import pwd
import re
import shutil
import sqlite3
import ssl
import stat
import subprocess
import sys
import time
import uuid
import urllib.request
import urllib.parse

ROOT = Path(__file__).resolve().parents[1]
UNIT_PATH = Path('/etc/systemd/system/rustproxy.service')
SHELL_PATH_PROFILE = Path('/etc/profile.d/rustproxy-path.sh')
SERVICE = 'rustproxy.service'
GEOIP_ASSETS = (
    'user-country-ipv4.csv',
    'user-country-ipv6.csv',
)
GEOIP_RELEASE = 'https://github.com/sapics/ip-location-db/releases/download/latest/'
GEOIP_CHECKSUM = 'https://github.com/sapics/ip-location-db/releases/download/checksum/'
GEOIP_MAX_BYTES = 256 * 1024 * 1024


def run(args: list[str], **kwargs) -> subprocess.CompletedProcess:
    return subprocess.run(args, check=True, stdin=subprocess.DEVNULL, **kwargs)


def safe_path(raw: str) -> Path:
    if not re.fullmatch(r'/[A-Za-z0-9_./-]+', raw) or '..' in raw.split('/') or raw in ('/', '/etc', '/var', '/opt', '/usr', '/var/lib'):
        raise ValueError('Use a dedicated absolute path with no spaces, traversal or shell characters.')
    p = Path(raw)
    for part in (p, *p.parents):
        if part.is_symlink():
            raise ValueError(f'Symlinked installation/data path is not allowed: {part}')
    return p


def address(raw: str, loopback: bool = False, allow_remote: bool = False) -> str:
    if raw.startswith('['):
        host, port = raw[1:].split(']:', 1)
    else:
        host, port = raw.rsplit(':', 1)
    ip = ipaddress.ip_address(host)
    if not port.isdecimal() or not 1 <= int(port) <= 65535:
        raise ValueError('Invalid listener IP/port.')
    if loopback and not ip.is_loopback and not allow_remote:
        raise ValueError('Administration must use loopback unless remote administration is explicitly accepted.')
    return f'[{ip}]:{int(port)}' if ip.version == 6 else f'{ip}:{int(port)}'


def listener_ip(raw: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address:
    host = raw[1:].split(']:', 1)[0] if raw.startswith('[') else raw.rsplit(':', 1)[0]
    return ipaddress.ip_address(host)




def admin_address(raw: str, allow_remote: bool = False) -> tuple[str, bool]:
    value = address(raw)
    ip = listener_ip(value)
    remote = not ip.is_loopback
    if remote and not allow_remote:
        raise ValueError('Remote administration requires --allow-remote-admin.')
    return value, remote

def digest(path: Path) -> str:
    with path.open('rb') as stream:
        result = hashlib.sha256()
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            result.update(block)
        return result.hexdigest()


def fetch_https(url: str, max_bytes: int) -> bytes:
    if urllib.parse.urlparse(url).scheme != 'https':
        raise ValueError('GeoIP downloads require HTTPS.')
    request = urllib.request.Request(url, headers={'User-Agent': 'RustProxy-installer/0.25.1'})
    with urllib.request.urlopen(request, timeout=30) as response:
        if urllib.parse.urlparse(response.geturl()).scheme != 'https':
            raise ValueError('GeoIP download redirected away from HTTPS.')
        length = response.headers.get('Content-Length')
        if length and int(length) > max_bytes:
            raise ValueError('GeoIP download is larger than the configured safety limit.')
        result = bytearray()
        while True:
            block = response.read(min(1024 * 1024, max_bytes + 1 - len(result)))
            if not block:
                break
            result.extend(block)
            if len(result) > max_bytes:
                raise ValueError('GeoIP download exceeded the configured safety limit.')
        return bytes(result)


def publisher_geoip_checksums() -> dict[str, str]:
    expected: dict[str, str] = {}
    for name in GEOIP_ASSETS:
        checksum = fetch_https(GEOIP_CHECKSUM + name + '.sha256', 4096).decode('ascii', 'strict').strip()
        matched = re.search(r'\b([0-9a-fA-F]{64})\b', checksum)
        if not matched:
            raise ValueError(f'Invalid publisher checksum for {name}.')
        expected[name] = matched.group(1).lower()
    return expected


def geoip_file_digest(path: Path) -> str | None:
    try:
        st = path.lstat()
    except FileNotFoundError:
        return None
    if not stat.S_ISREG(st.st_mode) or path.is_symlink():
        return None
    return digest(path)


def prepare_geoip_update(data: Path, prefix: Path) -> tuple[Path | None, list[str]]:
    """Verify local GeoIP files against current publisher checksums.

    Network/download work happens before the service is stopped. Changed files are
    staged under the code prefix and are copied into the private data transaction
    only after its backup exists.
    """
    expected = publisher_geoip_checksums()
    directory = data / 'geoip'
    if directory.exists() and directory.is_symlink():
        raise ValueError('GeoIP directory cannot be a symlink.')
    stale = [name for name in GEOIP_ASSETS if geoip_file_digest(directory / name) != expected[name]]
    if not stale:
        print('GeoIP database is current (publisher SHA-256 matches).')
        return None, []
    staging = prefix / ('.geoip-stage-' + uuid.uuid4().hex)
    staging.mkdir(mode=0o700)
    try:
        for name in stale:
            payload = fetch_https(GEOIP_RELEASE + name, GEOIP_MAX_BYTES)
            actual = hashlib.sha256(payload).hexdigest()
            if actual != expected[name]:
                raise ValueError(f'GeoIP SHA-256 mismatch for {name}.')
            write_private(staging / name, payload, 0o600)
        print('GeoIP update prepared: ' + ', '.join(stale))
        return staging, stale
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise


def copy_private(source: Path, target: Path, mode: int = 0o600) -> None:
    temporary = target.with_name(target.name + '.' + uuid.uuid4().hex + '.tmp')
    fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, mode)
    try:
        with source.open('rb') as input_stream, os.fdopen(fd, 'wb') as output_stream:
            shutil.copyfileobj(input_stream, output_stream, 1024 * 1024)
            output_stream.flush(); os.fsync(output_stream.fileno())
        os.replace(temporary, target)
        directory = os.open(target.parent, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
    finally:
        temporary.unlink(missing_ok=True)


def apply_geoip_update(data: Path, staging: Path | None, names: list[str], uid: int, gid: int) -> None:
    if not names:
        return
    directory = data / 'geoip'
    if directory.exists() and directory.is_symlink():
        raise ValueError('GeoIP directory cannot be a symlink.')
    directory.mkdir(mode=0o700, exist_ok=True)
    os.chmod(directory, 0o700)
    os.chown(directory, uid, gid)
    if staging is None:
        raise ValueError('GeoIP staging directory is missing.')
    for name in names:
        source = staging / name
        target = directory / name
        copy_private(source, target, 0o600)
        os.chown(target, uid, gid)
    print('GeoIP database updated: ' + ', '.join(names))


def download_geoip(data: Path, uid: int, gid: int) -> None:
    # First install has no code-prefix staging transaction, so download and verify
    # each asset directly before publishing it atomically into the private data dir.
    expected = publisher_geoip_checksums()
    directory = data / 'geoip'
    if directory.exists() and directory.is_symlink():
        raise ValueError('GeoIP directory cannot be a symlink.')
    directory.mkdir(mode=0o700, exist_ok=True)
    os.chmod(directory, 0o700)
    os.chown(directory, uid, gid)
    for name in GEOIP_ASSETS:
        payload = fetch_https(GEOIP_RELEASE + name, GEOIP_MAX_BYTES)
        actual = hashlib.sha256(payload).hexdigest()
        if actual != expected[name]:
            raise ValueError(f'GeoIP SHA-256 mismatch for {name}.')
        target = directory / name
        write_private(target, payload, 0o600)
        os.chown(target, uid, gid)


def write_private(path: Path, value: bytes, mode: int = 0o600) -> None:
    temporary = path.with_name(path.name + '.' + uuid.uuid4().hex + '.tmp')
    fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, mode)
    try:
        with os.fdopen(fd, 'wb') as stream:
            stream.write(value); stream.flush(); os.fsync(stream.fileno())
        os.replace(temporary, path)
        directory = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
    finally:
        temporary.unlink(missing_ok=True)


def shell_path_profile(prefix: Path) -> bytes:
    binary_dir = prefix / 'current'
    return (
        '# Managed by the RustProxy installer. New login shells can invoke rust-proxy directly.\n'
        'case ":${PATH:-}:" in\n'
        f'  *":{binary_dir}:"*) ;;\n'
        f'  *) export PATH="{binary_dir}${{PATH:+:$PATH}}" ;;\n'
        'esac\n'
    ).encode()


def install_shell_path(prefix: Path) -> None:
    write_private(SHELL_PATH_PROFILE, shell_path_profile(prefix), 0o644)


def check_shell_path(prefix: Path) -> bool:
    expected = shell_path_profile(prefix)
    try:
        st = SHELL_PATH_PROFILE.lstat()
        valid = stat.S_ISREG(st.st_mode) and st.st_uid == 0 and not (st.st_mode & 0o022) and SHELL_PATH_PROFILE.read_bytes() == expected
    except OSError:
        valid = False
    if not valid:
        print(
            f'WARNING: shell PATH setup is missing or outdated: {SHELL_PATH_PROFILE}. '
            f'Expected {prefix}/current in PATH; update mode only checks this file and does not modify it.',
            file=sys.stderr,
        )
    return valid


def read_record(prefix: Path) -> dict:
    path = prefix / 'install.json'
    fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW)
    with os.fdopen(fd, 'rb') as stream:
        st = os.fstat(stream.fileno())
        if not stat.S_ISREG(st.st_mode) or st.st_uid != 0 or st.st_mode & 0o077 or st.st_size > 16384:
            raise ValueError('Unsafe or invalid root-owned installation record.')
        value = json.load(stream)
    if value.get('schema') != 1 or value.get('prefix') != str(prefix):
        raise ValueError('Installation record mismatch.')
    safe_path(value['data_dir']); validate_user(value['user'])
    remote_admin = bool(value.get('remote_admin', False))
    admin, checked_remote = admin_address(value['admin'], remote_admin)
    if checked_remote != remote_admin:
        raise ValueError('Installation record remote administration flag does not match its address.')
    value['admin'] = admin
    value['remote_admin'] = remote_admin
    return value


def validate_user(name: str) -> None:
    if not re.fullmatch(r'[a-z_][a-z0-9_-]{0,30}', name) or name in ('root', 'nobody'):
        raise ValueError('Use a dedicated non-root service account.')


def service_unit(record: dict) -> str:
    prefix, data, user = record['prefix'], record['data_dir'], record['user']
    args = f'{prefix}/current/rust-proxy --data-dir {data} --admin-listen {record["admin"]} serve'
    args += f' --proxy-listen {record["http"]}'
    if record.get('https'):
        args += f' --tls-listen {record["https"]}'
    if record.get('remote_admin'):
        args += ' --allow-remote-admin'
    return f'''[Unit]
Description=RustProxy reverse proxy
Wants=network-online.target
After=network-online.target

[Service]
Type=simple
User={user}
Group={record['group']}
WorkingDirectory={data}
UMask=0077
Environment=RUST_LOG=rust_proxy=info,warn
Environment=MALLOC_ARENA_MAX=2
Environment=PATH=/opt/rustproxy-certbot/bin:/usr/local/bin:/usr/bin:/bin
EnvironmentFile=-/etc/rustproxy/environment
ExecStart={args}
Restart=on-failure
RestartSec=3
KillSignal=SIGINT
KillMode=control-group
TimeoutStopSec=120
AmbientCapabilities=CAP_NET_BIND_SERVICE
CapabilityBoundingSet=CAP_NET_BIND_SERVICE
NoNewPrivileges=true
PrivateTmp=true
PrivateDevices=true
ProtectSystem=strict
LogsDirectory=rustproxy
LogsDirectoryMode=0770
ReadWritePaths={data} /var/log/rustproxy
ProtectHome=true
ProtectKernelTunables=true
ProtectKernelModules=true
ProtectKernelLogs=true
ProtectControlGroups=true
RestrictSUIDSGID=true
RestrictNamespaces=true
RestrictRealtime=true
RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6
SystemCallArchitectures=native
LockPersonality=true

[Install]
WantedBy=multi-user.target
'''