#!/usr/bin/python3
"""Upgrade the router's package, vpn-router, or with --system every package of the system, recording the state and
the output.

Started by the Router Manager as the transient unit vpn-router-manager-upgrade, or with --system
vpn-router-manager-system-upgrade (systemd-run), so it runs outside the manager's unit: the
restarts of the manager and of the router's web service that the package's scripts make do not
stop it. Usage: upgrade [--system] <user>, the user who started it.

The state file, STATE_FILE or with --system SYSTEM_STATE_FILE, holds the state, running while it
runs, then succeeded or failed, with the user, the start and the end, the exit status of apt-get
and the installed versions of PACKAGES before and after; it is replaced by a rename, so a reader
never sees half of it. The log file, LOG_FILE or SYSTEM_LOG_FILE, holds the output of dpkg and
apt-get, from the start of the run. A dpkg run that was interrupted is finished first
(dpkg --configure -a). Then apt-get install --only-upgrade PACKAGES, or with --system apt-get
upgrade, which upgrades the installed packages without installing or removing any. apt-get waits
up to LOCK_SECONDS for the dpkg lock, which Ubuntu's apt-daily or unattended-upgrades may hold,
takes the installed version of a changed configuration file and asks nothing.
"""

import json
import os
import subprocess
import sys
import time

STATE_DIR = '/var/lib/vpn-router'
STATE_FILE = os.path.join(STATE_DIR, 'upgrade.json')
LOG_FILE = os.path.join(STATE_DIR, 'upgrade.log')
SYSTEM_STATE_FILE = os.path.join(STATE_DIR, 'system-upgrade.json')
SYSTEM_LOG_FILE = os.path.join(STATE_DIR, 'system-upgrade.log')

PACKAGES = ('vpn-router',)

#: Seconds apt-get waits for the dpkg lock.
LOCK_SECONDS = 300

#: The environment of dpkg and apt-get: no questions, the installed configuration files kept, and needrestart left out,
#: as the package's scripts restart their own services.
ENVIRONMENT = {
    'DEBIAN_FRONTEND': 'noninteractive',
    'UCF_FORCE_CONFFOLD': '1',
    'APT_LISTCHANGES_FRONTEND': 'none',
    'NEEDRESTART_SUSPEND': '1',
}


def now():
    return time.strftime('%Y-%m-%dT%H:%M:%S%z')


def versions():
    """The installed version of each of PACKAGES, None for one that is not installed."""
    found = {}
    for package in PACKAGES:
        result = subprocess.run(['dpkg-query', '-W', '-f=${Version}', package], capture_output=True, text=True)
        found[package] = (result.stdout.strip() or None) if result.returncode == 0 else None
    return found


def save(state, path):
    """Write the state to the file of path by a rename."""
    temporary = path + '.new'
    with open(temporary, 'w', encoding='utf-8') as handle:
        json.dump(state, handle)
        handle.write('\n')
    os.chmod(temporary, 0o644)
    os.replace(temporary, path)


def main(argv):
    system = argv[:1] == ['--system']
    if system:
        argv = argv[1:]
    if len(argv) != 1:
        print('usage: upgrade [--system] <user>', file=sys.stderr)
        return 2
    state_file, log_file = (SYSTEM_STATE_FILE, SYSTEM_LOG_FILE) if system else (STATE_FILE, LOG_FILE)
    os.makedirs(STATE_DIR, exist_ok=True)
    state = {'state': 'running', 'user': argv[0], 'started': now(), 'finished': None, 'exit_status': None,
             'before': versions(), 'after': None}
    save(state, state_file)
    environment = dict(os.environ, **ENVIRONMENT)
    with open(log_file, 'w', encoding='utf-8') as log:
        os.chmod(log_file, 0o644)
        status = 1
        for command in (
            ['dpkg', '--configure', '-a'],
            ['apt-get', '-y', '-q', '-o', f'DPkg::Lock::Timeout={LOCK_SECONDS}', '-o', 'Dpkg::Options::=--force-confdef',
             '-o', 'Dpkg::Options::=--force-confold', *(['upgrade'] if system else ['install', '--only-upgrade', *PACKAGES])],
        ):
            log.write(f'$ {" ".join(command)}\n')
            log.flush()
            try:
                status = subprocess.run(command, stdin=subprocess.DEVNULL, stdout=log, stderr=subprocess.STDOUT, env=environment).returncode
            except OSError as exc:
                log.write(f'{exc}\n')
                status = 1
            if status != 0:
                break
        log.write(f'exit status {status}\n')
    state.update(state='succeeded' if status == 0 else 'failed', finished=now(), exit_status=status, after=versions())
    save(state, state_file)
    return status


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))
