#!/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.24.8'})
    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
ProtectSystem=strict
LogsDirectory=rustproxy
LogsDirectoryMode=0770
ReadWritePaths={data} /var/log/rustproxy
ProtectHome=true
ProtectKernelTunables=true
ProtectKernelModules=true
ProtectControlGroups=true
RestrictSUIDSGID=true
RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6
LockPersonality=true

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


def checked_tree(source: Path) -> None:
    """Accept only regular files/directories and relative Certbot-internal symlinks."""
    boundary = source.resolve(strict=True)
    if source.is_symlink():
        raise ValueError('Data root cannot be a symlink.')
    for p in source.rglob('*'):
        st = p.lstat()
        if stat.S_ISLNK(st.st_mode):
            target = os.readlink(p)
            if os.path.isabs(target) or not p.resolve(strict=True).is_relative_to(boundary):
                raise ValueError('Data backup contains an external or absolute symlink.')
        elif not (stat.S_ISDIR(st.st_mode) or stat.S_ISREG(st.st_mode)):
            raise ValueError('Data backup contains a special file.')


def copy_data(source: Path, target: Path) -> None:
    checked_tree(source)
    shutil.copytree(source, target, symlinks=True)
    os.chmod(target, 0o700)


def chown_tree(path: Path, uid: int, gid: int) -> None:
    os.chown(path, uid, gid)
    for child in path.rglob('*'):
        os.chown(child, uid, gid, follow_symlinks=False)


def ensure_service_data_access(data: Path, account) -> None:
    """Repair the managed data-root metadata and verify service-user traversal.

    systemd applies WorkingDirectory after dropping to User=/Group=. A stale
    root-owned 0700 data directory therefore fails as status=200/CHDIR before
    RustProxy itself can emit diagnostics. Repair only the managed data root;
    never chmod/chown parent directories. Then fail before service downtime if a
    parent directory still prevents traversal.
    """
    try:
        st = data.lstat()
    except FileNotFoundError as error:
        raise ValueError(f'Data directory is missing: {data}') from error
    if stat.S_ISLNK(st.st_mode) or not stat.S_ISDIR(st.st_mode):
        raise ValueError(f'Data directory must be a real directory: {data}')

    os.chown(data, account.pw_uid, account.pw_gid)
    os.chmod(data, 0o700)

    probe = subprocess.run(
        ['runuser', '-u', account.pw_name, '--', 'test', '-x', str(data)],
        check=False, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    if probe.returncode != 0:
        raise ValueError(
            f'Service user {account.pw_name} cannot traverse data directory {data}; '
            'check execute/search permissions on its parent directories (for example with namei -l).'
        )


def switch(prefix: Path, release: str) -> None:
    if not re.fullmatch(r'[A-Za-z0-9_.-]+', release) or '..' in release:
        raise ValueError('Invalid release name.')
    temporary = prefix / ('current-' + uuid.uuid4().hex)
    os.symlink('releases/' + release, temporary)
    os.replace(temporary, prefix / 'current')


def running_release_matches(record: dict) -> bool:
    """Check the actual systemd process executable without publishing a version API.

    This installer targets Linux/systemd. Fail closed if /proc or MainPID cannot
    establish that the service runs the exact staged, version-checked binary.
    """
    try:
        result = subprocess.run(['systemctl', 'show', '--property=MainPID', '--value', SERVICE],
                                check=False, capture_output=True, text=True, timeout=3)
        if result.returncode != 0 or not result.stdout.strip().isdigit():
            return False
        pid = int(result.stdout.strip())
        if pid <= 1:
            return False
        expected = Path(record['prefix']) / 'releases' / record['release'] / 'rust-proxy'
        return Path(f'/proc/{pid}/exe').samefile(expected)
    except (OSError, ValueError, KeyError, subprocess.SubprocessError):
        return False


def admin_health_endpoint(record: dict) -> tuple[str, bool, str]:
    """Resolve the active persisted administration listener and scheme.

    The installer record contains only the startup CLI address. Since 0.9.x the
    dashboard can persist both a different admin bind and HTTPS in config.sqlite3.
    Update health checks must follow that active intent or a healthy HTTPS panel is
    mistaken for a failed release and rollback can fail for the same reason.
    """
    raw = record['admin']
    tls = False
    allowed_hosts: list[str] = []
    database = Path(record['data_dir']) / 'config.sqlite3'
    try:
        connection = sqlite3.connect(f'file:{database}?mode=ro', uri=True, timeout=2)
        try:
            row = connection.execute('SELECT document FROM configuration WHERE id = 1').fetchone()
        finally:
            connection.close()
        if row:
            document = json.loads(row[0])
            administration = document.get('administration') or {}
            if isinstance(administration, dict):
                if administration.get('listen'):
                    raw = address(str(administration['listen']))
                tls = bool(administration.get('tls', False))
                values = administration.get('allowed_hosts') or []
                if isinstance(values, list):
                    allowed_hosts = [str(value).strip() for value in values if str(value).strip()]
    except (OSError, sqlite3.Error, ValueError, TypeError, json.JSONDecodeError):
        # Old releases and damaged/uninitialized databases used the startup HTTP
        # listener. The real startup diagnostics remain authoritative if this
        # compatibility read cannot resolve persisted settings.
        pass

    host, port = (raw[1:].split(']:') if raw.startswith('[') else raw.rsplit(':', 1))
    ip = ipaddress.ip_address(host)
    connect_host = ('::1' if ip.version == 6 else '127.0.0.1') if ip.is_unspecified else host
    authority_host = allowed_hosts[0] if allowed_hosts else connect_host
    try:
        authority_ip = ipaddress.ip_address(authority_host)
        authority = f'[{authority_ip}]:{port}' if authority_ip.version == 6 else f'{authority_ip}:{port}'
    except ValueError:
        authority = f'{authority_host}:{port}'
    connect_raw = f'[{connect_host}]:{port}' if ':' in connect_host else f'{connect_host}:{port}'
    return connect_raw, tls, authority


def healthy(record: dict, version: str, attempts: int = 30) -> bool:
    raw, tls, health_host = admin_health_endpoint(record)
    host, port = (raw[1:].split(']:') if raw.startswith('[') else raw.rsplit(':', 1))
    for _ in range(attempts):
        if tls:
            # Local deployment health is process identity/liveness, not PKI validation.
            # The configured admin certificate may intentionally be self-signed or for
            # a DNS name while we connect to the local socket.
            conn = http.client.HTTPSConnection(host, int(port), timeout=2, context=ssl._create_unverified_context())
        else:
            conn = http.client.HTTPConnection(host, int(port), timeout=2)
        try:
            conn.request('GET', '/healthz', headers={'Host': health_host})
            response = conn.getresponse(); content = response.read(4097)
            active = subprocess.run(['systemctl', 'is-active', '--quiet', SERVICE], check=False, stdout=subprocess.DEVNULL).returncode == 0
            if (active and response.status == 200 and len(content) <= 4096
                    and json.loads(content).get('status') == 'ok' and record.get('version') == version
                    and running_release_matches(record)):
                return True
        except (OSError, ValueError, http.client.HTTPException, ssl.SSLError):
            pass
        finally:
            conn.close()
        time.sleep(1)
    return False


def stage(binary: Path, prefix: Path, user: str, sha: str) -> tuple[str, str]:
    release = 'build-' + sha[:24]
    releases = prefix / 'releases'; releases.mkdir(mode=0o755, exist_ok=True); os.chmod(releases, 0o755)
    directory = releases / release
    if not directory.exists():
        directory.mkdir(mode=0o755); os.chmod(directory, 0o755)
        try:
            shutil.copyfile(binary, directory / 'rust-proxy')
            os.chmod(directory / 'rust-proxy', 0o755)
        except BaseException:
            shutil.rmtree(directory); raise
    if directory.is_symlink() or (directory / 'rust-proxy').is_symlink() or digest(directory / 'rust-proxy') != sha:
        raise ValueError('Staged binary checksum mismatch.')
    # Never execute the supplied binary as root, even for its version banner.
    result = run(['runuser', '-u', user, '--', str(directory / 'rust-proxy'), '--version'], capture_output=True, text=True, timeout=15)
    matched = re.fullmatch(r'(?:rust-proxy|RustProxy) (\d+\.\d+\.\d+(?:[-+][\w.-]+)?)\s*', result.stdout)
    if not matched:
        raise ValueError('Binary did not return a RustProxy version.')
    return release, matched[1]


def backup_stopped(record: dict, prefix: Path) -> Path:
    # Service stop must happen before this call: ACME files and SQLite form one backup.
    data = Path(record['data_dir'])
    lock_fd = os.open(data / 'process.lock', os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW, 0o600)
    try:
        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        account = pwd.getpwnam(record['user']); os.fchown(lock_fd, account.pw_uid, account.pw_gid)
        backups = prefix / 'backups'; backups.mkdir(mode=0o700, exist_ok=True)
        backup = backups / ('before-' + record['version'] + '-' + uuid.uuid4().hex[:12]); backup.mkdir(mode=0o700)
        copy_data(data, backup / 'data')
        write_private(backup / 'install.json', json.dumps(record, indent=2).encode())
        if UNIT_PATH.exists():
            shutil.copyfile(UNIT_PATH, backup / 'rustproxy.service'); os.chmod(backup / 'rustproxy.service', 0o600)
        return backup
    finally:
        os.close(lock_fd)


def restore_service_unit(backup: Path) -> None:
    saved = backup / 'rustproxy.service'
    if saved.exists():
        copy_private(saved, UNIT_PATH, 0o644)
    else:
        UNIT_PATH.unlink(missing_ok=True)
    run(['systemctl', 'daemon-reload'])


def reset_saved_config(data: Path) -> None:
    """Development-only reset of desired config while preserving auth/certs/audit data."""
    database = data / 'config.sqlite3'
    if database.is_symlink() or not database.is_file():
        raise ValueError('Cannot reset configuration: config.sqlite3 is missing or unsafe.')
    connection = sqlite3.connect(f'file:{database}?mode=rw', uri=True, timeout=5)
    try:
        row = connection.execute("SELECT 1 FROM sqlite_master WHERE type='table' AND name='configuration'").fetchone()
        if row is None:
            raise ValueError('Cannot reset configuration: configuration table is missing.')
        connection.execute('DELETE FROM configuration WHERE id = 1')
        connection.commit()
    finally:
        connection.close()
    print('Stored vhost/settings configuration cleared; account, certificates and other private data were preserved.')


def service_failure_details() -> str:
    chunks: list[str] = []
    commands = (
        ['systemctl', 'show', SERVICE, '--property=ActiveState,SubState,ExecMainStatus', '--no-pager'],
        ['journalctl', '-u', SERVICE, '-n', '25', '--no-pager', '--output=cat'],
    )
    for command in commands:
        try:
            result = subprocess.run(command, check=False, capture_output=True, text=True, timeout=5)
            output = ((result.stdout or '') + (result.stderr or '')).strip()
            if output:
                chunks.append(output)
        except (OSError, subprocess.SubprocessError):
            pass
    return '\n'.join(chunks)[-8000:]


def stop_for_managed_update(data: Path) -> None:
    """Ask RustProxy 0.14.0+ to force hot cache into the enabled cold tier before exit.

    Older installed versions ignore the marker and still follow their ordinary graceful
    shutdown policy, so the first upgrade to 0.14.0 remains backward compatible.
    The marker is removed before backup so it never becomes persisted application data.
    """
    marker = data / '.managed-update'
    marker.unlink(missing_ok=True)
    write_private(marker, b'1\n', 0o644)
    try:
        run(['systemctl', 'stop', SERVICE])
    finally:
        marker.unlink(missing_ok=True)


def update(record: dict, binary: Path, sha: str, start: bool, refresh_geoip: bool = False, reset_config: bool = False) -> dict:
    prefix = Path(record['prefix']); data = Path(record['data_dir']); account = pwd.getpwnam(record['user'])
    # Updates deliberately do not rewrite shell environment. Check the installer-
    # managed PATH profile early and report the exact remediation target instead.
    check_shell_path(prefix)
    # Do this before staging/network work and, critically, before stopping the
    # currently healthy service. It also repairs stale root:root/0700 metadata
    # left by manual maintenance while preserving private 0700 semantics.
    ensure_service_data_access(data, account)
    release, version = stage(binary, prefix, record['user'], sha)
    geoip_stage: Path | None = None
    geoip_names: list[str] = []
    if refresh_geoip:
        geoip_stage, geoip_names = prepare_geoip_update(data, prefix)
    active = subprocess.run(['systemctl', 'is-active', '--quiet', SERVICE], check=False, stdout=subprocess.DEVNULL).returncode == 0
    stop_for_managed_update(data)
    try:
        backup = backup_stopped(record, prefix)
    except BaseException:
        if active: run(['systemctl', 'start', SERVICE])
        if geoip_stage is not None:
            shutil.rmtree(geoip_stage, ignore_errors=True)
        raise
    replacement = dict(record, release=release, version=version, sha256=sha)
    try:
        if reset_config:
            reset_saved_config(data)
        apply_geoip_update(data, geoip_stage, geoip_names, account.pw_uid, account.pw_gid)
        switch(prefix, release)
        write_private(UNIT_PATH, service_unit(replacement).encode(), 0o644)
        run(['systemctl', 'daemon-reload'])
        if start:
            run(['systemctl', 'start', SERVICE])
            if not healthy(replacement, version):
                raise RuntimeError('New release failed the local health/version check.')
        write_private(prefix / 'install.json', json.dumps(replacement, indent=2).encode())
    except BaseException as error:
        details = service_failure_details()
        if details:
            print('New release diagnostics:\n' + details, file=sys.stderr)
        stop_for_managed_update(data)
        failed = data.with_name(data.name + '.failed-' + uuid.uuid4().hex[:12])
        data.rename(failed)  # keep the failed generation for diagnosis; do not delete it
        copy_data(backup / 'data', data); chown_tree(data, account.pw_uid, account.pw_gid)
        switch(prefix, record['release'])
        restore_service_unit(backup)
        write_private(prefix / 'install.json', json.dumps(record, indent=2).encode())
        if active:
            run(['systemctl', 'start', SERVICE])
            if not healthy(record, record['version']):
                raise RuntimeError(f'Update failed: {error}. Rollback restored code/data/unit but previous service health also failed. Backup: {backup}; failed state: {failed}') from error
        raise RuntimeError(f'Update failed: {error}. Previous code, private data and service unit were restored. Installation is still present; do not run install.sh. Fix the cause and rerun update.sh. Backup: {backup}; failed state: {failed}') from error
    finally:
        if geoip_stage is not None:
            shutil.rmtree(geoip_stage, ignore_errors=True)
    print(f'Backup (contains private keys): {backup}')
    return replacement


def install(record: dict, binary: Path, sha: str, start: bool, install_geoip: bool = False) -> dict:
    import grp
    prefix = Path(record['prefix']); data = Path(record['data_dir'])
    if (prefix / 'install.json').exists() or (prefix / 'current').exists() or UNIT_PATH.exists():
        raise ValueError('Installation/service already exists. A failed update rolls back to the previous installed release, so rerun update.sh after fixing the reported cause. install.sh never overwrites an existing installation.')
    if data.exists() and any(data.iterdir()):
        raise ValueError('First installation requires an empty data directory; existing data will never be reset.')
    try:
        account = pwd.getpwnam(record['user'])
        if account.pw_uid == 0 or account.pw_shell not in ('/usr/sbin/nologin', '/sbin/nologin', '/bin/false'):
            raise ValueError('Existing service account must be non-root with a non-login shell.')
    except KeyError:
        run(['useradd', '--system', '--user-group', '--home-dir', str(data), '--no-create-home', '--shell', '/usr/sbin/nologin', record['user']])
        account = pwd.getpwnam(record['user'])
    record['group'] = grp.getgrgid(account.pw_gid).gr_name
    prefix.mkdir(parents=True, mode=0o755, exist_ok=True); os.chmod(prefix, 0o755)
    data.mkdir(parents=True, mode=0o700, exist_ok=True); os.chmod(data, 0o700); os.chown(data, account.pw_uid, account.pw_gid)
    if install_geoip:
        download_geoip(data, account.pw_uid, account.pw_gid)
    release, version = stage(binary, prefix, record['user'], sha)
    record.update(release=release, version=version, sha256=sha)
    switch(prefix, release)
    install_shell_path(prefix)
    # Record is written even if initial startup fails, so an update can repair it.
    write_private(prefix / 'install.json', json.dumps(record, indent=2).encode())
    write_private(UNIT_PATH, service_unit(record).encode(), 0o644)
    run(['systemctl', 'daemon-reload']); run(['systemctl', 'enable', SERVICE])
    if start:
        run(['systemctl', 'start', SERVICE])
        if not healthy(record, version):
            run(['systemctl', 'stop', SERVICE])
            raise RuntimeError('Installed but startup failed. Data was retained; inspect journalctl -u rustproxy.')
    return record


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('mode', choices=['install', 'update'])
    parser.add_argument('--binary', type=Path, default=ROOT / 'target/release/rust-proxy')
    parser.add_argument('--prefix', default='/opt/rustproxy')
    parser.add_argument('--data-dir', default='/var/lib/rustproxy')
    parser.add_argument('--user', default='rustproxy')
    parser.add_argument('--http', default='0.0.0.0:80')
    tls = parser.add_mutually_exclusive_group()
    tls.add_argument('--https', default='0.0.0.0:443', help='TLS address (default: 0.0.0.0:443); first certificate is self-signed')
    tls.add_argument('--no-https', action='store_true', help='Disable the default 443 TLS listener')
    parser.add_argument('--admin', default='127.0.0.1:9090')
    parser.add_argument('--allow-remote-admin', action='store_true', help='Allow --admin on 0.0.0.0/:: instead of loopback')
    parser.add_argument('--sha256', help='Expected SHA-256 of the supplied binary')
    parser.add_argument('--no-start', action='store_true', help='Leave the service stopped after install/update')
    parser.add_argument('--plan', action='store_true', help='Print a plan without root privileges, builds or host changes')
    parser.add_argument('--skip-geoip', action='store_true', help='Skip GeoIP install/update verification (country allow lists may use stale/missing data)')
    parser.add_argument('--reset-config', action='store_true', help='UPDATE ONLY, development recovery: discard saved vhosts/settings after backup while preserving account/certificates/private files')
    args = parser.parse_args()
    os.umask(0o077)
    try:
        prefix, data = safe_path(args.prefix), safe_path(args.data_dir)
        if args.reset_config and args.mode != 'update':
            raise ValueError('--reset-config is only valid with update.sh.')
        if prefix == data or prefix in data.parents or data in prefix.parents:
            raise ValueError('Code and private data must use separate non-nested directories.')
        validate_user(args.user)
        admin, remote_admin = admin_address(args.admin, args.allow_remote_admin)
        record = dict(schema=1, prefix=str(prefix), data_dir=str(data), user=args.user, http=address(args.http),
                      https=None if args.no_https else address(args.https), admin=admin, remote_admin=remote_admin)
        ports = [x.rsplit(':', 1)[1] for x in (record['http'], record['https'], record['admin']) if x]
        if len(set(ports)) != len(ports):
            raise ValueError('HTTP, TLS and administration ports must differ.')
        binary = args.binary.absolute()
        if binary.is_symlink() or not binary.is_file() or not os.access(binary, os.X_OK):
            raise ValueError('Build the release binary as an ordinary user first (bash dev.sh release); provide --binary.')
        sha = digest(binary)
        if args.sha256 and args.sha256.lower() != sha:
            raise ValueError('Binary SHA-256 mismatch; no changes made.')
        if args.plan:
            print(json.dumps({'operation':args.mode, 'binary':str(binary), 'sha256':sha, 'install_defaults':record,
                              'update_uses_saved_settings':args.mode == 'update', 'backup_before_update':True, 'rollback_code_and_data':True,
                              'start_service':not args.no_start,
                              'geoip_action':'skip' if args.skip_geoip else ('download-and-verify' if args.mode == 'install' else 'verify-and-refresh-if-stale'),
                              'shell_path_profile':str(SHELL_PATH_PROFILE),
                              'shell_path_action':'install' if args.mode == 'install' else 'check-only',
                              'reset_saved_config':bool(args.reset_config), 'changes_made':False}, indent=2)); return
        if sys.platform != 'linux' or os.geteuid() != 0 or not shutil.which('systemctl') or not shutil.which('runuser'):
            raise ValueError('Installation requires root on Linux with systemd and runuser. Use --plan for inspection.')
        prefix.mkdir(parents=True, mode=0o755, exist_ok=True)
        if prefix.stat().st_uid != 0 or prefix.stat().st_mode & 0o022:
            raise ValueError('Installation prefix must be root-owned and not group/world-writable.')
        fd = os.open(prefix / '.deploy.lock', os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW, 0o600)
        with os.fdopen(fd, 'w') as lock:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            result = install(record, binary, sha, not args.no_start, not args.skip_geoip) if args.mode == 'install' else update(
                read_record(prefix), binary, sha, not args.no_start, not args.skip_geoip, args.reset_config
            )
        print(f'RustProxy {result["version"]} installed. Administration: http://{result["admin"]}')
        print(f'Generate setup token: sudo -u {result["user"]} {prefix}/current/rust-proxy --data-dir {result["data_dir"]} admin setup-token')
        print(f'Initialize account: sudo -u {result["user"]} {prefix}/current/rust-proxy --data-dir {result["data_dir"]} admin init')
        if result.get('remote_admin'):
            print('WARNING: remote administration is plain HTTP. Restrict it with a firewall/private network; a protected tunnel remains safer.')
        else:
            print('Use an SSH tunnel for remote administration.')
        print('No demo backend or firewall changes were installed.')
    except (OSError, ValueError, RuntimeError, subprocess.SubprocessError, KeyError) as error:
        parser.exit(1, f'Deployment failed: {error}\n')


if __name__ == '__main__':
    main()