#!/usr/bin/python3
"""Stop the installed SVOS services for a VM backup, then restore their prior state."""
import argparse
import fcntl
import json
import os
from pathlib import Path
import stat
import subprocess
import sys
import tempfile
import time

SERVICES = ('svos-kernel.service', 'svos-interface.service', 'svos-terminal.service',
            'svos-knowledge.service', 'svos-postoffice.service', 'svos-voice.service')

class Quiesce:
    def __init__(self, directory, *, run=None):
        self.directory = Path(directory)
        self.directory.mkdir(mode=0o700, exist_ok=True)
        info = self.directory.lstat()
        if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.geteuid() or stat.S_IMODE(info.st_mode) != 0o700:
            raise RuntimeError('backup state directory must be private and owned by this user')
        self.state = self.directory / 'state.json'
        self.run = run or self._run

    @staticmethod
    def _run(args, **kwargs):
        return subprocess.run(['systemctl', *args], check=True, text=True, capture_output=True, **kwargs)

    def status(self, unit, deadline):
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError('service quiesce deadline exceeded')
        answer = self.run(['show', unit, '--property=ActiveState,Result'], timeout=min(5, remaining))
        fields = dict(line.split('=', 1) for line in answer.stdout.splitlines() if '=' in line)
        return fields.get('ActiveState'), fields.get('Result')

    def record(self, active):
        with tempfile.NamedTemporaryFile(mode='w', dir=self.directory, delete=False) as output:
            name = output.name
            json.dump({'version': 1, 'services': active}, output)
            output.flush()
            os.fsync(output.fileno())
        os.replace(name, self.state)
        self.sync_directory()

    def sync_directory(self):
        fd = os.open(self.directory, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)

    def freeze(self):
        if self.state.exists():
            raise RuntimeError('previous backup state remains; thaw before retrying freeze')
        deadline = time.monotonic() + 35
        active = []
        for unit in SERVICES:
            state, _ = self.status(unit, deadline)
            if state == 'active':
                active.append(unit)
            elif state not in ('inactive', 'failed'):
                raise RuntimeError('service is transitioning: ' + unit)
        self.record(active)
        try:
            if active:
                self.run(['stop', *active], timeout=max(0.1, deadline - time.monotonic()))
            for unit in active:
                state, result = self.status(unit, deadline)
                if state != 'inactive' or result != 'success':
                    raise RuntimeError('service did not stop cleanly: ' + unit)
        except BaseException:
            if active:
                try:
                    self.run(['start', '--no-block', *active], timeout=5)
                except Exception as error:
                    print('backup recovery start refused: ' + type(error).__name__, file=sys.stderr)
            # Keep the original state until thaw confirms recovery, even if freeze failed.
            raise

    def thaw(self):
        if not self.state.exists():
            return
        record = json.loads(self.state.read_text())
        active = record.get('services')
        if (set(record) != {'version', 'services'} or record['version'] != 1 or
                not isinstance(active, list) or any(unit not in SERVICES for unit in active) or
                len(active) != len(set(active))):
            raise RuntimeError('invalid backup recovery record')
        deadline = time.monotonic() + 15
        if active:
            self.run(['start', *active], timeout=10)
        for unit in active:
            if self.status(unit, deadline)[0] != 'active':
                raise RuntimeError('service did not resume: ' + unit)
        self.state.unlink()
        self.sync_directory()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=('freeze', 'thaw'))
    parser.add_argument('--state-dir', default='/run/svos-backup')
    args = parser.parse_args()
    try:
        hook = Quiesce(args.state_dir)
        fd = os.open(hook.directory / 'lock', os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            getattr(hook, args.action)()
        finally:
            os.close(fd)
    except Exception as error:
        print('SVOS backup ' + args.action + ' failed: ' + str(error), file=sys.stderr)
        return 1
    print('SVOS backup ' + args.action + ' complete')
    return 0

if __name__ == '__main__':
    raise SystemExit(main())