"""The backups of the router, for vpn-router-backup and the web service.

A backup is a consistent copy of the state that setup cannot make again, staged in STAGING: the
trees of TREES, the folders of the web sites and the accounts of the object storage among them, with backup.s3 its
buckets, the files of site_to_site psk_file outside them under files/,
the directory as slapcat writes it (ldap/vpn-router.ldif), with backup.metrics the database as pg_dump writes it
(postgresql/vpn-router.dump), the databases of the web sites and of the name server (DUMPS) once they exist, and
manifest.json. The trees and the directory are copied while
the configuration lock and the users lock are held, so the configuration, the PKI and the users
match; pg_dump reads a snapshot of its own.

backup.target local writes the staging directory as a tar archive compressed with zstd,
ARCHIVE, to backup.path, removes the archives of this host that the retention settings of the
schedule no longer keep, and with backup.rclone_remote syncs or copies the archives of this host
to the remote of RCLONE_CONFIG. backup.target pbs sends the staging directory to a Proxmox Backup
Server as the archive router.pxar of the snapshot host/<host>/<time>; the server keeps them.

A run holds RUN_LOCK and records itself in STATE: result running, ok, warning (the backup is
made, the cloud copy failed) or failed, with the times, the target, the archive or snapshot,
its size, the error and the cloud copy. restore() puts a backup back and runs setup.
"""

import contextlib
import datetime
import fcntl
import json
import os
import pathlib
import re
import shutil
import socket
import subprocess
import tempfile
import time

import vpnrouter
import vpnrouter_auth
import vpnrouter_dns
import vpnrouter_ldap
import vpnrouter_s3
import vpnrouter_sites
import vpnrouter_users

STATE = pathlib.Path('/var/lib/vpn-router/backup.json')
WORK_DIR = pathlib.Path('/var/lib/vpn-router/backup')
STAGING = WORK_DIR / 'staging'
RESTORE_DIR = WORK_DIR / 'restore'
#: What a restore replaces is kept here, in a directory named after the time of the restore.
BEFORE_RESTORE = WORK_DIR / 'before-restore'
RUN_LOCK = pathlib.Path('/run/vpn-router/backup.lock')
RCLONE_CONFIG = pathlib.Path('/etc/vpn-router/rclone.conf')
PBS_CLIENT = pathlib.Path('/usr/bin/proxmox-backup-client')
UNIT = 'vpn-router-backup.service'
TIMER = 'vpn-router-backup.timer'

#: The state a backup copies: directories and files, each where it is on the system; with backup.s3 also the buckets
#: of the object storage (vpnrouter_s3.DATA_DIR).
TREES = ('/etc/vpn-router', '/etc/wireguard', '/etc/ufw', '/etc/default/ufw', '/etc/vpn-router-manager',
         str(vpnrouter_sites.SITES_DIR), str(vpnrouter_s3.IAM_DIR))

#: The databases a backup always dumps, each while it exists: the dump -> the database and its password file.
DUMPS = {'postgresql/vpn-router-sites.dump': (vpnrouter_sites.DATABASE, vpnrouter_sites.DATABASE_SECRET),
         'postgresql/vpn-router-dns.dump': (vpnrouter_dns.DATABASE, vpnrouter_dns.DATABASE_SECRET)}

#: The name of a local archive: the host and the UTC time of the backup.
ARCHIVE = re.compile(r'^vpn-router-(?P<host>[A-Za-z0-9][A-Za-z0-9.-]*)-(?P<stamp>\d{8}T\d{6}Z)\.tar\.zst$')
STAMP = '%Y%m%dT%H%M%SZ'

LDIF = 'ldap/vpn-router.ldif'
DUMP = 'postgresql/vpn-router.dump'
PBS_ARCHIVE = 'router.pxar'

#: The period of each retention setting, as strftime formats of local time; weeks are ISO weeks.
PERIODS = {'keep_hourly': '%Y-%m-%d %H', 'keep_daily': '%Y-%m-%d', 'keep_weekly': '%G-%V', 'keep_monthly': '%Y-%m',
           'keep_yearly': '%Y'}

#: The units a restore stops while it replaces the directory and the database, and the ones it restarts after setup.
RESTORE_STOPS = ('vpn-router-web.service', 'vpn-router-metrics.service', 'vpn-router-webdav.service', 'pdns.service',
                 'vpn-router-s3.service', 'slapd.service')
RESTORE_RESTARTS = ('wg-quick@wg0.service', 'openvpn-server@vpn-router.service', 'strongswan.service',
                    'vpn-router-manager.service')


class Failure(Exception):
    """A backup or a restore that cannot be done."""


def now():
    """The local time in ISO 8601, to the second, with its offset."""
    return datetime.datetime.now().astimezone().isoformat(timespec='seconds')


def host():
    """The name of this host, which names its archives and its PBS backup group."""
    return socket.gethostname().split('.')[0]


def command(args, **kwargs):
    """Run a command; Failure with its standard error when it fails."""
    try:
        result = subprocess.run(args, capture_output=True, text=True, **kwargs)
    except OSError as exc:
        raise Failure(f'{args[0]}: {exc.strerror}')
    if result.returncode != 0:
        raise Failure(f'{" ".join(args[:2])} failed: {result.stderr.strip() or result.stdout.strip()}')
    return result.stdout


def write_state(state):
    """Write STATE in one step."""
    STATE.parent.mkdir(parents=True, exist_ok=True)
    descriptor, name = tempfile.mkstemp(dir=STATE.parent, prefix='.backup.json.')
    with os.fdopen(descriptor, 'w', encoding='utf-8') as handle:
        handle.write(json.dumps(state, indent=2) + '\n')
    os.chmod(name, 0o600)
    os.replace(name, STATE)


def running():
    """Whether a run holds RUN_LOCK."""
    try:
        with open(RUN_LOCK, encoding='utf-8') as handle:
            fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except FileNotFoundError:
        return False
    except BlockingIOError:
        return True
    return False


def state():
    """The state of the last run, {} when there was none; a run that ended without recording it is interrupted."""
    try:
        recorded = json.loads(STATE.read_text(encoding='utf-8'))
    except (OSError, ValueError):
        return {}
    if recorded.get('result') == 'running' and not running():
        recorded['result'] = 'interrupted'
    return recorded


@contextlib.contextmanager
def run_lock():
    """Hold RUN_LOCK; Failure when another run holds it."""
    RUN_LOCK.parent.mkdir(parents=True, exist_ok=True)
    with open(RUN_LOCK, 'a', encoding='utf-8') as handle:
        try:
            fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise Failure('another backup or restore is running')
        yield


def archives(path):
    """The archives of this host in path, newest first, as [{name, time, size}]; time is the UTC time of the backup."""
    found = []
    try:
        entries = list(os.scandir(path))
    except OSError:
        return found
    for entry in entries:
        match = ARCHIVE.match(entry.name)
        if match and match['host'] == host() and entry.is_file():
            when = datetime.datetime.strptime(match['stamp'], STAMP).replace(tzinfo=datetime.timezone.utc)
            found.append({'name': entry.name, 'time': when.isoformat(), 'size': entry.stat().st_size})
    return sorted(found, key=lambda item: item['time'], reverse=True)


def kept(found, keeps):
    """The names of the archives the retention settings keep, by the rules of Proxmox Backup Server's prune.

    found is newest first. The settings apply in the order of PERIODS: each keeps the newest archive of each of its
    last N periods that have archives and no archive a former setting kept, and the older archives of those periods
    are removed, so no later setting keeps them.
    """
    marks = {}
    for option, period in PERIODS.items():
        count = keeps.get(option, 0)
        if not count:
            continue
        periods = {moment(item).strftime(period) for item in found if marks.get(item['name']) == 'keep'}
        included = set()
        for item in found:
            if item['name'] in marks:
                continue
            key = moment(item).strftime(period)
            if key in periods:
                continue
            if key not in included:
                if len(included) >= count:
                    break
                included.add(key)
                marks[item['name']] = 'keep'
            else:
                marks[item['name']] = 'remove'
    return {name for name, mark in marks.items() if mark == 'keep'}


def moment(item):
    """The local time of an archive of archives()."""
    return datetime.datetime.fromisoformat(item['time']).astimezone()


def psk_files(cfg):
    """The psk_file of each site_to_site connection outside the trees of TREES."""
    files = []
    for connection in cfg['site_to_site'].values():
        name = connection['psk_file'].strip()
        if name and not any(name == tree or name.startswith(tree + '/') for tree in TREES) and os.path.isfile(name):
            files.append(name)
    return files


def database_args(cfg):
    """The command prefix and the database argument of pg_dump and pg_restore: postgres and vpn-router, or web.database."""
    configured = cfg['web']['database'].strip()
    if configured:
        return [], [f'--dbname={configured}']
    return ['runuser', '-u', 'postgres', '--'], [f'--dbname={vpnrouter_auth.DATABASE}']


def stage(cfg):
    """Copy what a backup holds to STAGING, emptied first, and return the manifest."""
    shutil.rmtree(STAGING, ignore_errors=True)
    STAGING.mkdir(parents=True, mode=0o700)
    trees = [tree for tree in (*TREES, *((str(vpnrouter_s3.DATA_DIR),) if cfg['backup']['s3'] else ()))
             if os.path.exists(tree)]
    files = psk_files(cfg)
    vpnrouter_users.LOCK_FILE.parent.mkdir(parents=True, exist_ok=True)
    with vpnrouter.config_lock(), open(vpnrouter_users.LOCK_FILE, 'w') as users:
        fcntl.flock(users, fcntl.LOCK_EX)
        command(['cp', '-a', '--parents', *trees, str(STAGING)])
        if files:
            (STAGING / 'files').mkdir()
            command(['cp', '-a', '--parents', *files, str(STAGING / 'files')])
        (STAGING / LDIF).parent.mkdir()
        command(['slapcat', '-b', vpnrouter_ldap.BASE, '-l', str(STAGING / LDIF)])
    items = trees + [f'files{name}' for name in files] + [LDIF]
    dumps = {**({DUMP: database_args(cfg)} if cfg['backup']['metrics'] else {}),
             **{item: (['runuser', '-u', 'postgres', '--'], [f'--dbname={name}'])
                for item, (name, secret) in DUMPS.items() if secret.is_file()}}
    for item, (prefix, database) in dumps.items():
        (STAGING / item).parent.mkdir(exist_ok=True)
        with open(STAGING / item, 'wb') as dump:
            result = subprocess.run([*prefix, 'pg_dump', '-Fc', *database], stdout=dump, stderr=subprocess.PIPE)
        if result.returncode != 0:
            raise Failure(f'pg_dump for {item} failed: {result.stderr.decode(errors="replace").strip()}')
        items.append(item)
    version = subprocess.run(['dpkg-query', '-W', '-f=${Version}', 'vpn-router'], capture_output=True, text=True)
    manifest = {'host': host(), 'time': now(), 'version': version.stdout.strip() or None,
                'schedule': cfg['backup']['schedule'], 'items': items}
    (STAGING / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n', encoding='utf-8')
    return manifest


def pbs_environment(cfg):
    """The environment of proxmox-backup-client: the repository, the token secret and the fingerprint of backup."""
    if not PBS_CLIENT.exists():
        raise Failure(f'{PBS_CLIENT} does not exist; install proxmox-backup-client-static')
    backup = cfg['backup']
    environment = dict(os.environ, PBS_REPOSITORY=backup['pbs_repository'], PBS_PASSWORD=backup['pbs_password'])
    if backup['pbs_fingerprint']:
        environment['PBS_FINGERPRINT'] = backup['pbs_fingerprint']
    return environment


def pbs_options(cfg):
    """The namespace and key file options of proxmox-backup-client from backup."""
    options = []
    if cfg['backup']['pbs_namespace']:
        options += ['--ns', cfg['backup']['pbs_namespace']]
    if cfg['backup']['pbs_keyfile']:
        options += ['--keyfile', cfg['backup']['pbs_keyfile']]
    return options


def cloud(cfg):
    """Sync or copy the archives of this host in backup.path to backup.rclone_remote; Failure when it fails."""
    backup = cfg['backup']
    rclone = shutil.which('rclone')
    if not rclone:
        raise Failure('rclone is not installed')
    if not RCLONE_CONFIG.is_file():
        raise Failure(f'{RCLONE_CONFIG} does not exist; make the remote with rclone config --config {RCLONE_CONFIG}')
    command([rclone, backup['rclone_mode'], backup['path'], backup['rclone_remote'], '--config', str(RCLONE_CONFIG),
             '--include', f'vpn-router-{host()}-*.tar.zst'])


def run():
    """Make a backup to backup.target as the configuration says, recording it in STATE; the recorded state."""
    cfg = vpnrouter.load_config()
    backup = cfg['backup']
    started = time.time()
    record = {'result': 'running', 'started': now(), 'finished': None, 'target': backup['target'], 'archive': None,
              'snapshot': None, 'size': None, 'removed': [], 'error': None, 'cloud': None}
    with run_lock():
        write_state(record)
        try:
            stage(cfg)
            if backup['target'] == 'pbs':
                epoch = int(started)
                command([str(PBS_CLIENT), 'backup', f'{PBS_ARCHIVE}:{STAGING}', '--backup-type', 'host',
                         '--backup-id', host(), '--backup-time', str(epoch), *pbs_options(cfg)],
                        env=pbs_environment(cfg))
                stamp = datetime.datetime.fromtimestamp(epoch, datetime.timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
                record['snapshot'] = f'host/{host()}/{stamp}'
                record['size'] = sum(path.stat().st_size for path in STAGING.rglob('*') if path.is_file())
            else:
                path = pathlib.Path(backup['path'])
                path.mkdir(parents=True, mode=0o700, exist_ok=True)
                name = f'vpn-router-{host()}-{time.strftime(STAMP, time.gmtime(started))}.tar.zst'
                partial = path / f'{name}.partial'
                command(['tar', '--zstd', '-cf', str(partial), '-C', str(STAGING), '.'])
                partial.chmod(0o600)
                partial.rename(path / name)
                record['archive'] = name
                record['size'] = (path / name).stat().st_size
                found = archives(path)
                keep = kept(found, {option: backup[option] for option in vpnrouter.backup_keeps(backup['schedule'])})
                for item in found:
                    if item['name'] not in keep:
                        (path / item['name']).unlink()
                        record['removed'].append(item['name'])
                if backup['rclone_remote']:
                    record['cloud'] = {'remote': backup['rclone_remote'], 'mode': backup['rclone_mode'],
                                       'result': 'ok', 'error': None}
                    try:
                        cloud(cfg)
                    except Failure as exc:
                        record['cloud'].update(result='failed', error=str(exc))
            record['result'] = 'warning' if record['cloud'] and record['cloud']['result'] != 'ok' else 'ok'
        except (Failure, OSError) as exc:
            record['result'] = 'failed'
            record['error'] = str(exc)
        finally:
            shutil.rmtree(STAGING, ignore_errors=True)
            record['finished'] = now()
            write_state(record)
    return record


def snapshots(cfg):
    """The PBS snapshots of this host, as proxmox-backup-client lists them."""
    namespace = cfg['backup']['pbs_namespace']
    return command([str(PBS_CLIENT), 'snapshot', 'list', f'host/{host()}', *(['--ns', namespace] if namespace else [])],
                   env=pbs_environment(cfg))


def extract(cfg, source, pbs):
    """Unpack a local archive, or with pbs fetch a snapshot, to RESTORE_DIR, emptied first; the manifest."""
    shutil.rmtree(RESTORE_DIR, ignore_errors=True)
    RESTORE_DIR.mkdir(parents=True, mode=0o700)
    if pbs:
        command([str(PBS_CLIENT), 'restore', source, PBS_ARCHIVE, str(RESTORE_DIR), *pbs_options(cfg)],
                env=pbs_environment(cfg))
    else:
        command(['tar', '--zstd', '-xf', source, '-C', str(RESTORE_DIR)])
    try:
        return json.loads((RESTORE_DIR / 'manifest.json').read_text(encoding='utf-8'))
    except (OSError, ValueError):
        raise Failure(f'{source} is not a backup of vpn-router-backup: it has no manifest.json')


def put_back(relative, target, saved):
    """Move target to saved, keeping its path, and copy the staged relative path of RESTORE_DIR in its place."""
    target = pathlib.Path(target)
    if target.exists() or target.is_symlink():
        keep = saved / target.relative_to('/')
        keep.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(target, keep)
    target.parent.mkdir(parents=True, exist_ok=True)
    command(['cp', '-a', str(RESTORE_DIR / relative), str(target)])


def restore(source, pbs, confirm):
    """Put the backup back: a local archive, or with pbs a snapshot; confirm(manifest) is asked before any change.

    What is replaced is kept in BEFORE_RESTORE. The services of the directory and the database are stopped while they
    are replaced, and setup applies the configuration afterwards. Returns the directory of what was replaced, or None
    when confirm declined.
    """
    cfg = vpnrouter.load_config()
    with run_lock():
        try:
            manifest = extract(cfg, source, pbs)
            if not confirm(manifest):
                return None
            saved = BEFORE_RESTORE / time.strftime(STAMP, time.gmtime())
            saved.mkdir(parents=True, mode=0o700)
            for unit in RESTORE_STOPS:
                subprocess.run(['systemctl', 'stop', unit], capture_output=True)
            for item in manifest['items']:
                if item.startswith('/'):
                    put_back(item.lstrip('/'), item, saved)
                elif item.startswith('files/'):
                    put_back(item, item[len('files'):], saved)
            if LDIF in manifest['items']:
                directory = vpnrouter_ldap.DB_DIRECTORY
                if directory.exists():
                    shutil.move(directory, saved / 'ldap')
                directory.mkdir(mode=0o700)
                command(['slapadd', '-b', vpnrouter_ldap.BASE, '-l', str(RESTORE_DIR / LDIF)])
                command(['chown', '-R', 'openldap:openldap', str(directory)])
            subprocess.run(['systemctl', 'start', 'slapd.service'], capture_output=True)
            restored = vpnrouter.load_config()
            if not restored['web']['database'].strip():
                vpnrouter_auth.reset_password(vpnrouter_auth.DATABASE, vpnrouter_auth.DATABASE_SECRET)
                vpnrouter_auth.provision(restored)
            dumps = {DUMP: database_args(restored)}
            for item, (name, secret) in DUMPS.items():
                if item in manifest['items']:
                    vpnrouter_auth.reset_password(name, secret)
                    vpnrouter_auth.provision_database(name, secret)
                    dumps[item] = (['runuser', '-u', 'postgres', '--'], [f'--dbname={name}'])
            for item, (prefix, database) in dumps.items():
                if item not in manifest['items']:
                    continue
                with open(RESTORE_DIR / item, 'rb') as dump:
                    result = subprocess.run([*prefix, 'pg_restore', '--clean', '--if-exists', *database], stdin=dump,
                                            stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
                if result.returncode != 0:
                    raise Failure(f'pg_restore of {item} failed: {result.stderr.decode(errors="replace").strip()}')
            subprocess.run(['systemctl', 'start', 'vpn-router-web.service'], capture_output=True)
            command(['systemctl', 'restart', 'vpn-router-setup.service'])
            for unit in RESTORE_RESTARTS:
                subprocess.run(['systemctl', 'try-restart', unit], capture_output=True)
            subprocess.run(['ufw', 'reload'], capture_output=True)
            return saved
        finally:
            shutil.rmtree(RESTORE_DIR, ignore_errors=True)