#!/usr/bin/env python3
"""Read-only offsite restore rehearsal. Only marked scratch data may be deleted.

The shell launcher holds the shared backup lock through this process. Docker has
one bind mount, no network/ports, no restart policy and no persistent log driver.
"""
import argparse
import datetime as dt
from email.message import EmailMessage
from email.policy import SMTP
import hashlib
import json
import logging
from logging.handlers import RotatingFileHandler
import os
from pathlib import Path
import re
import selectors
import shutil
import signal
import socket
import subprocess
import sys
import tarfile
import time

import backup_set

GIB = 1024 ** 3
MARKER = 'app-backup-restore-test-v1\n'
DOCKER = ['docker', '--host=unix:///var/run/docker.sock']
OWNER_LABEL = 'de.jr-it-services.app-backup.restore-test'


class TestError(Exception):
    pass


def atomic_json(path, data):
    temp = path.with_suffix('.tmp')
    temp.write_text(json.dumps(data, ensure_ascii=True, indent=2) + '\n')
    temp.replace(path)


def private_root(workdir, create=True):
    work = Path(workdir)
    # Do not follow aliases to source paths or accept broad filesystem roots.
    if (not work.is_absolute() or work.resolve() != work or len(work.parts) < 3 or
            any(c in str(work) for c in ':\n\r')):
        raise TestError('Unsafe WORKDIR for restore test')
    if work.exists() and (work.stat().st_uid != 0 or work.stat().st_mode & 0o022):
        raise TestError('WORKDIR must be root-owned and not writable by group/others')
    root = work / 'restore-test'
    if root.is_symlink():
        raise TestError('Restore-test root must not be a symlink')
    if not root.exists():
        if not create:
            return root
        root.mkdir(mode=0o700, parents=True)
        (root / '.owner').write_text(MARKER)
    marker = root / '.owner'
    if (root.stat().st_uid != 0 or root.stat().st_mode & 0o077 or
            not marker.is_file() or marker.is_symlink() or marker.read_text() != MARKER):
        raise TestError('Unrecognized/non-private restore-test directory; refusing cleanup')
    # These paths are written by us only. Never follow replacements into live data.
    for name in ('.owner', 'scratch', 'state.json', 'state.tmp', 'last-success.json',
                 'last-success.tmp', 'report.eml', 'restore-test.log', 'restore-test.log.1', 'db-created',
                 'child.json', 'child.tmp'):
        if (root / name).is_symlink():
            raise TestError('Symlink inside restore-test control directory')
    return root


class Rehearsal:
    def __init__(self, args):
        self.args = args
        self.root = private_root(args.workdir, create=args.mode != 'cleanup')
        self.scratch = self.root / 'scratch'
        self.owner = hashlib.sha256(str(self.root).encode()).hexdigest()[:16]
        self.container = 'app-backup-restore-test-' + self.owner
        self.started = time.time()
        self.deadline = time.monotonic() + args.timeout
        self.state = dict(status='RUNNING', started_epoch=int(self.started), run='',
                          components=[], cleanup='PENDING', error='')
        self.log = logging.getLogger('restore-test')
        self.log.setLevel(logging.INFO)
        for handler in self.log.handlers:
            handler.close()
        self.log.handlers.clear()
        if self.root.exists():
            handler = RotatingFileHandler(self.root / 'restore-test.log', maxBytes=512*1024, backupCount=1)
            handler.setFormatter(logging.Formatter('%(asctime)s %(message)s'))
            self.log.addHandler(handler)
            self.log.addHandler(logging.StreamHandler(sys.stdout))
        self.next_usage = 0
        self.used = 0
        self.monitor = False

    def remember_child(self, proc):
        stat = Path(f'/proc/{proc.pid}/stat').read_text().rsplit(')', 1)[1].split()
        atomic_json(self.root / 'child.json', dict(pid=proc.pid, start=stat[19]))

    def stop_orphan(self):
        marker = self.root / 'child.json'
        if not marker.exists():
            return
        child = json.loads(marker.read_text())
        pid = child['pid']
        if not isinstance(pid, int) or pid <= 1:
            raise TestError('Invalid child-process marker')
        def alive():
            try:
                stat = Path(f'/proc/{pid}/stat').read_text().rsplit(')', 1)[1].split()
                return stat[19] == child['start'] and stat[0] != 'Z' and os.getpgid(pid) == pid
            except (FileNotFoundError, ProcessLookupError):
                return False
        if alive():
            os.killpg(pid, signal.SIGTERM)
            end = time.monotonic() + 5
            while alive() and time.monotonic() < end:
                time.sleep(0.1)
            if alive():
                os.killpg(pid, signal.SIGKILL)
                end = time.monotonic() + 5
                while alive() and time.monotonic() < end:
                    time.sleep(0.1)
                if alive():
                    raise TestError('Child process could not be stopped; cleanup deferred')
        marker.unlink()

    def save(self):
        atomic_json(self.root / 'state.json', self.state)

    def disk(self, additional=0, force=False):
        if not self.monitor:
            return
        if time.monotonic() >= self.deadline:
            raise TestError('Restore-test time limit exceeded')
        free = shutil.disk_usage(self.root).free
        if free < self.args.min_free_gb * GIB + additional:
            raise TestError('Free-space reserve reached; test stopped and scratch will be cleaned')
        if force or time.monotonic() >= self.next_usage:
            self.used = 0
            for path, directories, files in os.walk(self.scratch, followlinks=False):
                for name in directories + files:
                    self.used += (Path(path) / name).lstat().st_blocks * 512
            self.next_usage = time.monotonic() + 2
        if self.used + additional > self.args.max_work_gb * GIB:
            raise TestError('Restore-test disk budget exceeded')

    def command(self, command, *, capture=0, output=None, limit=None, stdin=None, timeout=120):
        """Bound stdout and runtime, monitor disk, reap children even on interruption.

        stderr is intentionally not copied: DB errors can contain restored data.
        Errors identify the failed operation without logging SQL or credentials.
        """
        result = bytearray()
        total = 0
        end = time.monotonic() + timeout
        proc = subprocess.Popen(command, stdin=stdin or subprocess.DEVNULL,
                                stdout=subprocess.PIPE if capture or output is not None else subprocess.DEVNULL,
                                stderr=subprocess.DEVNULL, start_new_session=True)
        selector = selectors.DefaultSelector()
        if proc.stdout:
            selector.register(proc.stdout, selectors.EVENT_READ)
        try:
            self.remember_child(proc)
            while selector.get_map() or proc.poll() is None:
                self.disk()
                if time.monotonic() >= end:
                    raise TestError('Timed out: ' + Path(command[0]).name)
                if not selector.get_map():
                    time.sleep(0.1)
                    continue
                for key, _ in selector.select(0.2):
                    chunk = os.read(key.fd, 64*1024)
                    if not chunk:
                        selector.unregister(key.fileobj)
                        continue
                    total += len(chunk)
                    if total > (limit if limit is not None else capture):
                        raise TestError('Remote output exceeds expected size')
                    if output is not None:
                        self.disk(additional=len(chunk))
                        output.write(chunk)
                    else:
                        result.extend(chunk)
            if proc.wait() != 0:
                raise TestError('Command failed: ' + Path(command[0]).name)
            return bytes(result)
        finally:
            selector.close()
            if proc.poll() is None:
                os.killpg(proc.pid, signal.SIGTERM)
                try:
                    proc.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    os.killpg(proc.pid, signal.SIGKILL)
                    proc.wait()
            if proc.stdout:
                proc.stdout.close()
            (self.root / 'child.json').unlink(missing_ok=True)

    def remote(self, *args, **kwargs):
        return self.command(['bash', self.args.remote_check, self.args.config, *args],
                            timeout=self.args.timeout, **kwargs)

    def cleanup(self):
        self.stop_orphan()
        # db-created is written BEFORE docker run. A lost CLI reply cannot orphan a DB.
        if (self.root / 'db-created').exists():
            names = self.command(DOCKER + ['container', 'ls', '-a', '--filter',
                                          f'label={OWNER_LABEL}={self.owner}', '--format', '{{.Names}}'],
                                 capture=4096).decode().splitlines()
            if any(name != self.container for name in names):
                raise TestError('Unexpected container ownership; refusing cleanup')
            if names:
                self.command(DOCKER + ['container', 'rm', '-f', '-v', self.container])
            (self.root / 'db-created').unlink()
        if self.scratch.exists():
            # shutil.rmtree uses fd-relative traversal on Linux; nested symlinks
            # are unlinked, never traversed. No broad cleanup of WORKDIR or Docker.
            if self.scratch.is_symlink() or self.scratch.is_mount():
                raise TestError('Unsafe scratch directory; refusing cleanup')
            shutil.rmtree(self.scratch)
        self.state['cleanup'] = 'OK'

    def preflight(self):
        self.monitor = True
        for executable in ('tar', 'gzip', 'zstd'):
            if not shutil.which(executable):
                raise TestError('Missing tool: ' + executable)
        if any(label.startswith('db-') for label in self.args.expected):
            self.check_image()
        self.disk(force=True)

    def check_image(self):
        if hasattr(self, 'image_id'):
            return
        if not self.args.image or self.args.image.startswith('-') or any(c.isspace() for c in self.args.image):
            raise TestError('Set RESTORE_TEST_DB_IMAGE to a locally prepared MariaDB image matching production')
        raw = self.command(DOCKER + ['image', 'inspect', self.args.image], capture=128*1024)
        info = json.loads(raw)[0]
        volumes = info.get('Config', {}).get('Volumes') or {}
        if set(volumes) - {'/var/lib/mysql'}:
            raise TestError('Test image declares unexpected volumes')
        # Store immutable local image ID: a tag change during a test cannot switch binaries.
        self.image_id = info['Id']

    def select(self):
        base = self.args.remote.rstrip('/')
        if self.args.run:
            runs = [self.args.run]
        else:
            listing = self.remote('lsf', base, '--dirs-only', '--max-depth', '1', capture=1024*1024).decode()
            runs = sorted((line.rstrip('/') for line in listing.splitlines()
                           if backup_set.RUN.fullmatch(line.rstrip('/')) and
                           line.startswith(self.args.prefix + '_')), reverse=True)
        for run in runs:
            if not backup_set.RUN.fullmatch(run) or not run.startswith(self.args.prefix + '_'):
                raise TestError('Invalid remote run name')
            files = self.remote('lsf', base + '/' + run, '--files-only', '--max-depth', '1', capture=1024*1024).decode().splitlines()
            if 'COMPLETE' not in files:
                if self.args.run:
                    raise TestError('Selected backup is incomplete')
                self.log.info('Incomplete backup skipped: %s', run)
                continue
            self.state['run'] = run
            self.save()
            for name, maximum in [('manifest.json', 1024*1024), ('COMPLETE', 256)]:
                data = self.remote('cat', base + '/' + run + '/' + name, capture=maximum)
                (self.scratch / name).write_bytes(data)
            manifest = backup_set.load_set(self.scratch, metadata_only=True)
            if manifest['run_id'] != run:
                raise TestError('Remote directory and manifest run differ')
            age = time.time() - manifest['created_epoch']
            self.state['backup_age_hours'] = round(age / 3600, 2)
            if age < -300 or age > self.args.max_age_hours * 3600:
                raise TestError('Backup is stale or dated in the future')
            labels = {item['label'] for item in manifest['archives']}
            if set(self.args.expected) - labels:
                raise TestError('Backup is missing configured components')
            allowed = {item['file'] for item in manifest['archives']} | {'manifest.json', 'COMPLETE', 'UPLOADED'}
            if set(files) != allowed and set(files) != allowed - {'UPLOADED'}:
                raise TestError('Remote backup has missing or unexpected files')
            return manifest
        raise TestError('No completed remote backup found')

    def expand_size(self, archive):
        # Full tar validation is performed first by the shared validator. A second
        # streaming pass estimates logical size, including sparse files and metadata.
        unpacked = 0
        count = 0
        proc = subprocess.Popen(['zstd' if archive.suffix == '.zst' else 'gzip', '-dc', '--', str(archive)],
                                stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, start_new_session=True)
        try:
            self.remember_child(proc)
            with tarfile.open(fileobj=proc.stdout, mode='r|') as tar:
                for item in tar:
                    self.disk()
                    unpacked += ((item.size + 4095) // 4096) * 4096 + 4096
                    count += 1
            while proc.stdout.read(1024*1024):
                self.disk()
            if proc.wait() != 0:
                raise TestError('Archive decompression failed')
            return unpacked, count
        finally:
            proc.stdout.close()
            if proc.poll() is None:
                proc.terminate()
            proc.wait()
            (self.root / 'child.json').unlink(missing_ok=True)

    def database(self, sql):
        self.check_image()
        # No imported code can reach a host DB socket, production path or network.
        data = self.scratch / 'database'
        data.mkdir(mode=0o700)
        (self.root / 'db-created').write_text(self.container + '\n')
        startup = '''set -eu
mkdir -p /var/lib/mysql/tmp
mariadb-install-db --no-defaults --datadir=/var/lib/mysql/db --auth-root-authentication-method=normal --skip-test-db --user=root --tmpdir=/var/lib/mysql/tmp >/dev/null 2>&1
exec mariadbd --no-defaults --user=root --datadir=/var/lib/mysql/db --socket=/var/lib/mysql/mysql.sock --pid-file=/var/lib/mysql/mysql.pid --tmpdir=/var/lib/mysql/tmp --skip-networking --skip-log-bin --event-scheduler=OFF --local-infile=0 --innodb-buffer-pool-size=64M --innodb-log-file-size=16M --max-allowed-packet=64M
'''
        self.command(DOCKER + ['run', '-d', '--pull=never', '--name', self.container,
                     '--label', f'{OWNER_LABEL}={self.owner}', '--network=none', '--read-only',
                     '--cap-drop=ALL', '--security-opt=no-new-privileges', '--user=0:0',
                     '--memory=512m', '--memory-swap=512m', '--cpus=1', '--pids-limit=128',
                     '--log-driver=none', '--no-healthcheck', '--volume', str(data) + ':/var/lib/mysql:Z',
                     '--entrypoint=/bin/sh', self.image_id, '-c', startup])
        client = DOCKER + ['exec', '-i', self.container, 'mariadb', '--no-defaults',
                          '--socket=/var/lib/mysql/mysql.sock', '-uroot', '--batch', '--skip-column-names']
        ready = time.monotonic() + 90
        while True:
            try:
                self.command(client + ['--execute=SELECT 1'], capture=64, timeout=10)
                break
            except TestError:
                self.disk()
                if time.monotonic() > ready:
                    raise TestError('Isolated MariaDB did not become ready') from None
                time.sleep(0.5)
        self.command(client + ['--execute=CREATE DATABASE restore_test'])
        with sql.open('rb') as stream:
            # --binary-mode disables client commands in batch SQL (e.g. system/source).
            self.command(client + ['--binary-mode', 'restore_test'], stdin=stream, timeout=self.args.timeout)
        self.command(DOCKER + ['exec', self.container, 'mariadb-check', '--no-defaults',
                              '--socket=/var/lib/mysql/mysql.sock', '-uroot', '--all-databases'], timeout=self.args.timeout)
        tables = self.command(client + ["--execute=SELECT COUNT(*) FROM information_schema.tables WHERE table_schema NOT IN ('mysql','information_schema','performance_schema','sys')"], capture=128).decode().strip()
        return int(tables)

    def execute(self):
        self.preflight()
        manifest = self.select()
        self.log.info('Testing %s, %s components, sequentially', self.state['run'], len(manifest['archives']))
        for item in manifest['archives']:
            component = dict(label=item['label'], status='RUNNING', archive_bytes=item['size'])
            self.state['components'].append(component)
            self.save()
            self.log.info('Component %s: download, verify, extract', item['label'])
            self.disk(additional=item['size'], force=True)
            archive = self.scratch / item['file']
            with archive.open('wb') as output:
                self.remote('cat', self.args.remote.rstrip('/') + '/' + self.state['run'] + '/' + item['file'],
                            output=output, limit=item['size'])
            if archive.stat().st_size != item['size'] or backup_set.digest(archive) != item['sha256']:
                raise TestError('Archive checksum/size mismatch: ' + item['label'])
            # Run shared validation as a monitored child, so even corrupt streams
            # and metadata-heavy archives cannot bypass the total timeout.
            self.command([sys.executable, str(Path(__file__).with_name('backup_set.py')),
                          'tar', str(archive), item['root']], timeout=self.args.timeout)
            unpacked, entries = self.expand_size(archive)
            self.disk(additional=unpacked, force=True)
            extracted = self.scratch / 'extracted'
            extracted.mkdir(mode=0o700)
            self.command(['tar', '--extract', '--file', str(archive), '--acls', '--xattrs', '--selinux',
                          '--numeric-owner', '--same-owner', '-C', str(extracted)], timeout=self.args.timeout)
            # The compressed input is no longer needed; free it BEFORE DB import.
            archive.unlink()
            component.update(extracted_bytes=unpacked, entries=entries)
            if item['label'].startswith('db-'):
                sql = extracted / item['root']
                if not sql.is_file() or sql.is_symlink() or not sql.stat().st_size:
                    raise TestError('SQL dump missing or empty')
                self.disk(additional=512*1024*1024, force=True)
                component['tables'] = self.database(sql)
            self.monitor = False
            self.cleanup()
            self.scratch.mkdir(mode=0o700)
            self.monitor = True
            component['status'] = 'SUCCESS'
            self.save()
            self.log.info('Component %s: SUCCESS; files and test DB removed', item['label'])
        self.state['status'] = 'SUCCESS'

    def report(self):
        fqdn = socket.getfqdn()
        host = self.args.hostname or fqdn
        state = self.state
        message = EmailMessage(policy=SMTP)
        message['From'] = self.args.sender or 'app-backup@' + fqdn
        message['To'] = self.args.to
        message['Subject'] = f"[app-restore-test] {state['status']} {host} {state['run']}"
        message['Date'] = dt.datetime.now().astimezone()
        body = ['Restore-Test', f'Host: {host}', f"Status: {state['status']}", f"Sicherung: {state['run']}",
                f"Alter: {state.get('backup_age_hours', '?')} Stunden", f"Dauer: {state['duration_seconds']}s",
                f"Bereinigung: {state['cleanup']}", f"Fehler: {state['error']}", '', 'Komponenten:']
        for c in state['components']:
            body.append(f"{c['label']}: {c['status']}; Archiv {c['archive_bytes']} Bytes; "
                        f"Einträge {c.get('entries', '?')}; Tabellen {c.get('tables', '-')}")
        body.extend(['', 'Dateien entpackt; enthaltene DB-Dumps isoliert importiert und Tabellen geprüft.',
                     'Dies ist kein Start-/Funktionstest der produktiven Anwendungen.'])
        message.set_content('\n'.join(body))
        raw = message.as_bytes()
        (self.root / 'report.eml').write_bytes(raw)
        if self.args.mail == 'true':
            if not self.args.to.strip() or not str(message['From']).strip():
                raise TestError('Empty mail recipient or sender')
            with io_bytes(raw) as stream:
                self.command([self.args.sendmail, '-oi', '-t'], stdin=stream, timeout=self.args.mail_timeout)

    def finish(self):
        self.monitor = False
        for sig in (signal.SIGINT, signal.SIGTERM):
            signal.signal(sig, signal.SIG_IGN)
        try:
            self.cleanup()
        except Exception as exc:
            self.state.update(status='FAIL', cleanup='FAILED', error='Cleanup failed: ' + str(exc))
            self.log.error('%s', self.state['error'])
        for c in self.state['components']:
            if c['status'] == 'RUNNING':
                c['status'] = 'FAIL'
        self.state['duration_seconds'] = max(0, int(time.time() - self.state['started_epoch']))
        self.state['finished_epoch'] = int(time.time())
        self.save()
        if self.state['status'] == 'SUCCESS':
            atomic_json(self.root / 'last-success.json', self.state)
        try:
            self.report()
            self.state['mail'] = 'SENT' if self.args.mail == 'true' else 'DISABLED'
        except Exception as exc:
            self.state['mail'] = 'FAILED'
            self.log.error('Restore-test report failed (%s); inspect state.json', type(exc).__name__)
        self.save()
        self.log.info('Restore-test %s; cleanup %s', self.state['status'], self.state['cleanup'])
        return 0 if self.state['status'] == 'SUCCESS' and self.state['mail'] != 'FAILED' else 1


def io_bytes(raw):
    # A real temporary FD for Popen stdin, unlinked immediately on creation.
    import tempfile
    stream = tempfile.TemporaryFile()
    stream.write(raw)
    stream.seek(0)
    return stream


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', choices=('run', 'check', 'cleanup'))
    for key in ('workdir', 'prefix', 'remote', 'run', 'image', 'hostname', 'to', 'sender', 'sendmail', 'mail'):
        parser.add_argument('--' + key, required=True)
    for key in ('remote-check', 'config'):
        parser.add_argument('--' + key)
    parser.add_argument('--expected', action='append', default=[])
    for key in ('max-age-hours', 'timeout', 'mail-timeout'):
        parser.add_argument('--' + key, type=int, required=True)
    for key in ('max-work-gb', 'min-free-gb'):
        parser.add_argument('--' + key, type=float, required=True)
    args = parser.parse_args()
    if not (0 < args.max_work_gb <= 100000 and 0 <= args.min_free_gb <= 100000 and
            1 <= args.max_age_hours <= 8760 and 1 <= args.timeout <= 86400 and 1 <= args.mail_timeout <= 300):
        parser.error('Invalid restore-test limits')
    if os.geteuid() != 0:
        parser.error('Run as root via app-restore-test.sh')
    test = Rehearsal(args)
    if args.mode == 'cleanup' and not test.root.exists():
        return 0
    if args.mode == 'check':
        if any((test.root / name).exists() for name in ('scratch', 'db-created', 'child.json')):
            raise TestError('Interrupted test data found; run --cleanup before --check')
        test.preflight()
        test.log.info('Restore-test prerequisites and disk reserve OK; no image pulled or backup downloaded')
        return 0
    previous = test.root / 'state.json'
    if previous.exists():
        try:
            old = json.loads(previous.read_text())
            if not isinstance(old, dict) or not {'status', 'cleanup', 'started_epoch', 'components', 'run'} <= old.keys():
                raise ValueError('Invalid state')
        except (ValueError, OSError):
            old = dict(test.state)
        if old.get('status') == 'RUNNING' or old.get('cleanup') != 'OK':
            test.state = old
            test.state.update(status='FAIL', error='Previous restore test interrupted; recovery cleanup')
            result = test.finish()
            if args.mode == 'cleanup' or test.state['cleanup'] != 'OK':
                return result
            test = Rehearsal(args)
    if args.mode == 'cleanup':
        test.cleanup()
        return 0
    def interrupted(sig, _):
        raise TestError('Interrupted by signal ' + str(sig))
    for sig in (signal.SIGINT, signal.SIGTERM):
        signal.signal(sig, interrupted)
    try:
        test.cleanup()
        test.state['cleanup'] = 'PENDING'
        test.save()
        test.scratch.mkdir(mode=0o700)
        test.execute()
    except Exception as exc:
        test.state.update(status='FAIL', error=str(exc))
        test.log.error('Restore-test failed: %s', exc)
    return test.finish()


if __name__ == '__main__':
    os.umask(0o077)
    try:
        sys.exit(main())
    except Exception as exc:
        print(f'ERROR: restore-test {type(exc).__name__}: {exc}', file=sys.stderr)
        sys.exit(1)