#!/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 directly from the
sapics/ip-location-db GitHub release and validates their structure before publish.
Updates download/validate current files before downtime, compare content SHA-256
locally and stage only changed assets. 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[2]
UNIT_PATH = Path('/etc/systemd/system/rustproxy.service')
SHELL_PATH_PROFILE = Path('/etc/profile.d/rustproxy-path.sh')
SERVICE = 'rustproxy.service'
GEOIP_SOURCES = {
    'user-country-ipv4.csv': 'https://github.com/sapics/ip-location-db/releases/download/latest/user-country-ipv4.csv',
    'user-country-ipv6.csv': 'https://github.com/sapics/ip-location-db/releases/download/latest/user-country-ipv6.csv',
}
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.28.0'})
    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 acquire_geoip_update_lock(data: Path, uid: int | None = None, gid: int | None = None) -> int:
    """Serialize installer/manual GeoIP refresh with the in-process worker."""
    path = data / '.geoip-update.lock'
    flags = os.O_RDWR | os.O_CREAT
    if hasattr(os, 'O_NOFOLLOW'):
        flags |= os.O_NOFOLLOW
    fd = os.open(path, flags, 0o600)
    try:
        os.fchmod(fd, 0o600)
        if uid is not None and gid is not None:
            os.fchown(fd, uid, gid)
        fcntl.flock(fd, fcntl.LOCK_EX)
        return fd
    except BaseException:
        os.close(fd)
        raise


def release_geoip_update_lock(fd: int | None) -> None:
    if fd is None:
        return
    try:
        fcntl.flock(fd, fcntl.LOCK_UN)
    finally:
        os.close(fd)


def validate_geoip_csv(name: str, payload: bytes) -> None:
    """Reject truncated/error payloads before a GeoIP CSV is published."""
    family = 4 if name == 'user-country-ipv4.csv' else 6
    try:
        text = payload.decode('utf-8', 'strict')
    except UnicodeDecodeError as error:
        raise ValueError(f'GeoIP {name} is not valid UTF-8.') from error
    previous_end: int | None = None
    rows = 0
    for number, raw in enumerate(text.splitlines(), 1):
        line = raw.strip()
        if not line or line.startswith('#'):
            continue
        fields = [value.strip().strip('"') for value in line.split(',')]
        if len(fields) != 3:
            raise ValueError(f'GeoIP {name} row {number} must contain exactly 3 columns.')
        try:
            start = ipaddress.ip_address(fields[0])
            end = ipaddress.ip_address(fields[1])
        except ValueError as error:
            raise ValueError(f'GeoIP {name} row {number} contains an invalid IP address.') from error
        if start.version != family or end.version != family or int(start) > int(end):
            raise ValueError(f'GeoIP {name} row {number} contains an invalid address range.')
        if previous_end is not None and int(start) <= previous_end:
            raise ValueError(f'GeoIP {name} ranges overlap or are unsorted at row {number}.')
        country = fields[2]
        if country and not re.fullmatch(r'[A-Za-z]{2}', country):
            raise ValueError(f'GeoIP {name} row {number} contains an invalid country code.')
        previous_end = int(end)
        rows += 1
    if rows == 0:
        raise ValueError(f'GeoIP {name} does not contain any address ranges.')


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]]:
    """Download and validate current GeoIP CSVs before service downtime.

    Publisher checksums are intentionally not required here: the upstream checksum
    release can move independently from the `latest` data release. Local SHA-256 is
    used only to decide whether the validated payload changes an installed file.
    """
    directory = data / 'geoip'
    if directory.exists() and directory.is_symlink():
        raise ValueError('GeoIP directory cannot be a symlink.')
    payloads: dict[str, bytes] = {}
    stale: list[str] = []
    for name, url in GEOIP_SOURCES.items():
        payload = fetch_https(url, GEOIP_MAX_BYTES)
        validate_geoip_csv(name, payload)
        payloads[name] = payload
        if geoip_file_digest(directory / name) != hashlib.sha256(payload).hexdigest():
            stale.append(name)
    if not stale:
        print('GeoIP database is current (validated content is unchanged).')
        return None, []
    staging = prefix / ('.geoip-stage-' + uuid.uuid4().hex)
    staging.mkdir(mode=0o700)
    try:
        for name in stale:
            write_private(staging / name, payloads[name], 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. Validate each direct
    # payload before publishing it atomically into the private data directory.
    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, url in GEOIP_SOURCES.items():
        payload = fetch_https(url, GEOIP_MAX_BYTES)
        validate_geoip_csv(name, payload)
        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=/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
'''