#!/usr/bin/env python3
"""Install or validate the fixed Truf host-agent deployment."""

import json
import os
from pathlib import Path
import stat
import subprocess
import sys


PROJECT = Path('/opt/truf')
DEPLOY = PROJECT / 'deploy/host-agent'
INSTALL_ROOT = Path('/usr/lib/truf-host-agent')
SYSTEMD = Path('/etc/systemd/system')
TMPFILES = Path('/etc/tmpfiles.d/truf-host-agent.conf')
ACTIVE = Path('/etc/truf/runtime')
WORKER_PACKAGES = Path('/etc/truf/worker-packages')
DEPLOYMENT_PROFILE = Path('/etc/truf/deployment-profile')
AGENT_SOCKET = Path('/run/truf/host-agent.sock')
RUNTIME_UID = RUNTIME_GID = 10001
RUNTIME_GROUP = 'truf-runtime'
MAX_COPY_BYTES = 4 * 1024 * 1024
MAX_COMPOSE_OUTPUT_BYTES = 4 * 1024 * 1024
UNITS = ('truf-host-agent.socket', 'truf-host-agent.service')
SCRIPTS = ('truf_host_agent.py', 'truf_host_agent_install.py')
FIXED_ENV = {
    'PATH': '/usr/sbin:/usr/bin:/sbin:/bin',
    'LANG': 'C',
    'LC_ALL': 'C',
}
STANDALONE_PROFILE = {
    'name': 'standalone-edge-v1',
    'compose_files': ('compose.yaml', 'compose.edge.yaml'),
    'runtime_network': None,
    'runtime_ports': [{
        'mode': 'host', 'target': 443, 'published': '443', 'protocol': 'tcp',
    }],
    'runtime_cpus': 2.0,
    'runtime_mem_limit': str(6 * 1024 ** 3),
    'edge_cap_add': ['NET_BIND_SERVICE'],
    'data_volume': {'name': 'truf-docker_data'},
}
SHARED_HOST_PROFILE = {
    'name': 'shared-host-edge-v1',
    'compose_files': ('compose.yaml', 'compose.shared-host.yaml'),
    'runtime_network': 'host',
    'runtime_ports': None,
    'runtime_cpus': 0.9,
    'runtime_mem_limit': str(720 * 1024 ** 2),
    'edge_cap_add': None,
    'data_volume': {'name': 'truf-remote-server-data', 'external': True},
}


def _compose_config_command(profile):
    return (
        '/usr/bin/docker', 'compose', '--ansi', 'never', '--project-name',
        'truf-docker', '--env-file', '/etc/truf-edge/edge.env',
        '--project-directory', '/opt/truf',
        *(item for name in profile['compose_files'] for item in (
            '--file', '/opt/truf/' + name,
        )),
        'config', '--format', 'json',
    )


COMPOSE_CONFIG_COMMAND = _compose_config_command(STANDALONE_PROFILE)
EXPECTED_RUNTIME_MOUNTS = (
    ('volume', 'data', '/data', False, None),
    ('bind', '/etc/truf/runtime', '/data/config', True, False),
    ('bind', '/etc/truf/worker-packages', '/data/worker-packages', True, False),
    (
        'bind', '/var/lib/truf/runtime-document-candidates',
        '/data/runtime-document-candidates', False, False,
    ),
    ('bind', '/run/truf/host-agent.sock', '/run/truf/host-agent.sock', True, False),
    (
        'bind', '/var/lib/truf/host-agent/results',
        '/data/host-agent-results', True, False,
    ),
    ('bind', '/run/truf-postgres', '/run/truf-postgres', False, False),
)


class InstallError(RuntimeError):
    def __init__(self, category):
        self.category = str(category)
        super().__init__('host-agent deployment validation failed')


def _deployment_profile():
    descriptor = None
    try:
        descriptor = os.open(
            DEPLOYMENT_PROFILE,
            os.O_RDONLY | getattr(os, 'O_CLOEXEC', 0)
            | getattr(os, 'O_NOFOLLOW', 0),
        )
        before = os.fstat(descriptor)
        if (
            not stat.S_ISREG(before.st_mode) or before.st_uid != 0
            or before.st_gid != 0 or stat.S_IMODE(before.st_mode) != 0o444
            or before.st_nlink != 1 or before.st_size > 64
        ):
            raise InstallError('profile')
        payload = os.read(descriptor, 65)
        after = os.fstat(descriptor)
        if (
            len(payload) > 64 or before.st_dev != after.st_dev
            or before.st_ino != after.st_ino or before.st_mode != after.st_mode
            or before.st_uid != after.st_uid or before.st_gid != after.st_gid
            or before.st_nlink != after.st_nlink or before.st_size != after.st_size
            or before.st_mtime_ns != after.st_mtime_ns
        ):
            raise InstallError('profile')
    except FileNotFoundError:
        return STANDALONE_PROFILE
    except InstallError:
        raise
    except OSError:
        raise InstallError('profile') from None
    finally:
        if descriptor is not None:
            os.close(descriptor)
    try:
        name = payload.decode('ascii').strip()
    except UnicodeDecodeError:
        raise InstallError('profile') from None
    profiles = {
        STANDALONE_PROFILE['name']: STANDALONE_PROFILE,
        SHARED_HOST_PROFILE['name']: SHARED_HOST_PROFILE,
    }
    if name not in profiles:
        raise InstallError('profile')
    return profiles[name]


def _run(command, timeout=120):
    try:
        subprocess.run(
            tuple(command), stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL, cwd='/', env=FIXED_ENV, shell=False,
            timeout=timeout, check=True,
        )
    except Exception:
        raise InstallError('command') from None


def _capture(command, timeout=120):
    try:
        result = subprocess.run(
            tuple(command), stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL, cwd='/', env=FIXED_ENV, shell=False,
            timeout=timeout, check=True,
        )
        if len(result.stdout) > MAX_COMPOSE_OUTPUT_BYTES:
            raise InstallError('command')
        return result.stdout
    except InstallError:
        raise
    except Exception:
        raise InstallError('command') from None


def _unit_active(unit):
    if unit not in UNITS:
        raise InstallError('command')
    try:
        result = subprocess.run(
            ('/usr/bin/systemctl', 'is-active', '--quiet', unit),
            stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL, cwd='/', env=FIXED_ENV, shell=False,
            timeout=120, check=False,
        )
    except Exception:
        raise InstallError('command') from None
    if result.returncode not in (0, 3, 4):
        raise InstallError('command')
    return result.returncode == 0


def _validate_compose_projection(payload, profile=None):
    profile = profile or STANDALONE_PROFILE
    try:
        projection = json.loads(payload)
        if type(projection) is not dict or projection.get('name') != 'truf-docker':
            raise InstallError('compose')
        volume_definitions = projection.get('volumes')
        if (
            type(volume_definitions) is not dict
            or volume_definitions.get('data') != profile['data_volume']
        ):
            raise InstallError('compose')
        services = projection.get('services')
        runtime = services.get('runtime') if type(services) is dict else None
        edge = services.get('edge') if type(services) is dict else None
        if (
            type(runtime) is not dict or type(edge) is not dict
            or runtime.get('network_mode') != profile['runtime_network']
            or runtime.get('ports') != profile['runtime_ports']
            or runtime.get('cpus') != profile['runtime_cpus']
            or runtime.get('mem_limit') != profile['runtime_mem_limit']
            or edge.get('network_mode') != 'service:runtime'
            or edge.get('cap_add') != profile['edge_cap_add']
            or edge.get('image') != 'truf-local:edge'
        ):
            raise InstallError('compose')
        mounts = runtime.get('volumes') if type(runtime) is dict else None
        if type(mounts) is not list:
            raise InstallError('compose')
        observed = []
        for mount in mounts:
            if type(mount) is not dict or type(mount.get('read_only', False)) is not bool:
                raise InstallError('compose')
            kind = mount.get('type')
            if kind == 'volume':
                if set(mount) - {'type', 'source', 'target', 'read_only'}:
                    raise InstallError('compose')
                create_host_path = None
            elif kind == 'bind':
                if set(mount) - {'type', 'source', 'target', 'read_only', 'bind'}:
                    raise InstallError('compose')
                binding = mount.get('bind')
                if binding not in ({}, {'create_host_path': False}):
                    raise InstallError('compose')
                create_host_path = False
            else:
                raise InstallError('compose')
            observed.append((
                kind, mount.get('source'), mount.get('target'),
                mount.get('read_only', False), create_host_path,
            ))
        if tuple(observed) != EXPECTED_RUNTIME_MOUNTS:
            raise InstallError('compose')
    except InstallError:
        raise
    except Exception:
        raise InstallError('compose') from None


def _details(path, *, directory, uid, gid, mode):
    try:
        value = os.stat(path, follow_symlinks=False)
    except OSError:
        raise InstallError('metadata') from None
    expected = stat.S_ISDIR if directory else stat.S_ISREG
    if (
        not expected(value.st_mode) or value.st_uid != uid or value.st_gid != gid
        or stat.S_IMODE(value.st_mode) != mode
        or (not directory and value.st_nlink != 1)
    ):
        raise InstallError('metadata')
    return value


def _read_source(path, maximum=MAX_COPY_BYTES):
    descriptor = None
    try:
        descriptor = os.open(
            path, os.O_RDONLY | getattr(os, 'O_CLOEXEC', 0)
            | getattr(os, 'O_NOFOLLOW', 0),
        )
        before = os.fstat(descriptor)
        if (
            not stat.S_ISREG(before.st_mode) or before.st_nlink != 1
            or before.st_uid != 0 or stat.S_IMODE(before.st_mode) & 0o022
        ):
            raise InstallError('source')
        with os.fdopen(descriptor, 'rb') as handle:
            descriptor = None
            payload = handle.read(maximum + 1)
            after = os.fstat(handle.fileno())
        if len(payload) > maximum or (before.st_dev, before.st_ino, before.st_size) != (
            after.st_dev, after.st_ino, after.st_size,
        ):
            raise InstallError('source')
        return payload
    except InstallError:
        raise
    except Exception:
        raise InstallError('source') from None
    finally:
        if descriptor is not None:
            os.close(descriptor)


def _secure_tree(path):
    try:
        root = os.stat(path, follow_symlinks=False)
        if (
            not stat.S_ISDIR(root.st_mode)
            or root.st_uid != 0
            or stat.S_IMODE(root.st_mode) & 0o022
        ):
            raise InstallError('project')
        for current, directories, files in os.walk(path, topdown=True, followlinks=False):
            current_path = Path(current)
            entries = ((name, True) for name in directories)
            entries = tuple(entries) + tuple((name, False) for name in files)
            current_details = os.stat(current_path, follow_symlinks=False)
            if (
                not stat.S_ISDIR(current_details.st_mode)
                or current_details.st_uid != 0
                or stat.S_IMODE(current_details.st_mode) & 0o022
            ):
                raise InstallError('project')
            for name, directory in entries:
                details = os.stat(current_path / name, follow_symlinks=False)
                expected = stat.S_ISDIR if directory else stat.S_ISREG
                if (
                    not expected(details.st_mode)
                    or details.st_uid != 0
                    or stat.S_IMODE(details.st_mode) & 0o022
                    or (not directory and details.st_nlink != 1)
                ):
                    raise InstallError('project')
    except InstallError:
        raise
    except Exception:
        raise InstallError('project') from None


def _same_file(installed, source):
    if not _read_source(installed) == _read_source(source):
        raise InstallError('installed_content')


def _ensure_install_root():
    try:
        INSTALL_ROOT.mkdir(mode=0o755)
    except FileExistsError:
        pass
    except OSError:
        raise InstallError('write') from None
    _details(INSTALL_ROOT, directory=True, uid=0, gid=0, mode=0o755)


def _root_directory(path):
    try:
        details = os.stat(path, follow_symlinks=False)
    except OSError:
        raise InstallError('metadata') from None
    if (
        not stat.S_ISDIR(details.st_mode)
        or details.st_uid != 0
        or stat.S_IMODE(details.st_mode) & 0o022
    ):
        raise InstallError('metadata')


def _secure_executable(path):
    try:
        link = os.lstat(path)
        resolved = os.path.realpath(path)
        target = os.stat(path)
        parent = os.stat(Path(path).parent, follow_symlinks=False)
    except OSError:
        raise InstallError('executable') from None
    if (
        link.st_uid != 0
        or not (stat.S_ISREG(link.st_mode) or stat.S_ISLNK(link.st_mode))
        or not resolved.startswith(('/usr/bin/', '/usr/sbin/'))
        or not stat.S_ISREG(target.st_mode) or target.st_uid != 0
        or stat.S_IMODE(target.st_mode) & 0o022
        or not stat.S_ISDIR(parent.st_mode) or parent.st_uid != 0
        or stat.S_IMODE(parent.st_mode) & 0o022
    ):
        raise InstallError('executable')


def _runtime_group_exists():
    if sys.platform != 'linux':
        raise InstallError('group')
    try:
        import grp
        named = grp.getgrnam(RUNTIME_GROUP)
        numbered = grp.getgrgid(RUNTIME_GID)
    except KeyError:
        return False
    except Exception:
        raise InstallError('group') from None
    if named.gr_gid != RUNTIME_GID or numbered.gr_name != RUNTIME_GROUP:
        raise InstallError('group')
    return True


def _ensure_runtime_group():
    if _runtime_group_exists():
        return
    try:
        import grp
        grp.getgrgid(RUNTIME_GID)
    except KeyError:
        pass
    except Exception:
        raise InstallError('group') from None
    else:
        raise InstallError('group')
    _secure_executable('/usr/sbin/groupadd')
    _run(('/usr/sbin/groupadd', '--system', '--gid', str(RUNTIME_GID), RUNTIME_GROUP))
    if not _runtime_group_exists():
        raise InstallError('group')


def _validate_agent_socket():
    try:
        details = os.stat(AGENT_SOCKET, follow_symlinks=False)
    except OSError:
        raise InstallError('socket') from None
    if (
        not stat.S_ISSOCK(details.st_mode)
        or details.st_uid != 0
        or details.st_gid != RUNTIME_GID
        or stat.S_IMODE(details.st_mode) != 0o660
    ):
        raise InstallError('socket')


def _write(path, payload, *, uid, gid, mode, replace):
    temporary = path.parent / ('.' + path.name + '.truf-install')
    descriptor = None
    try:
        try:
            os.unlink(temporary)
        except FileNotFoundError:
            pass
        descriptor = os.open(
            temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL
            | getattr(os, 'O_NOFOLLOW', 0), mode,
        )
        os.fchmod(descriptor, mode)
        os.fchown(descriptor, uid, gid)
        view = memoryview(payload)
        while view:
            written = os.write(descriptor, view)
            if written <= 0:
                raise OSError('short write')
            view = view[written:]
        os.fsync(descriptor)
        os.close(descriptor)
        descriptor = None
        if not replace and path.exists():
            os.unlink(temporary)
            return
        os.replace(temporary, path)
        parent = os.open(path.parent, os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0))
        try:
            os.fsync(parent)
        finally:
            os.close(parent)
    except Exception:
        try:
            os.unlink(temporary)
        except OSError:
            pass
        raise InstallError('write') from None
    finally:
        if descriptor is not None:
            os.close(descriptor)
        payload = None


def validate(*, require_socket=True):
    if sys.platform != 'linux' or not hasattr(os, 'geteuid') or os.geteuid() != 0:
        raise InstallError('root')
    _root_directory(PROJECT.parent)
    _root_directory(INSTALL_ROOT.parent)
    if not _runtime_group_exists():
        raise InstallError('group')
    _details(PROJECT, directory=True, uid=0, gid=0, mode=0o755)
    _details(DEPLOY, directory=True, uid=0, gid=0, mode=0o755)
    _details(INSTALL_ROOT, directory=True, uid=0, gid=0, mode=0o755)
    _secure_tree(PROJECT / 'app')
    _details(ACTIVE, directory=True, uid=0, gid=0, mode=0o755)
    _details(WORKER_PACKAGES, directory=True, uid=0, gid=0, mode=0o755)
    _details(ACTIVE / 'config.yaml', directory=False, uid=RUNTIME_UID, gid=RUNTIME_GID, mode=0o600)
    _details(ACTIVE / 'secrets.yaml', directory=False, uid=RUNTIME_UID, gid=RUNTIME_GID, mode=0o600)
    layouts = (
        ('/var/lib/truf/runtime-document-candidates', RUNTIME_UID, RUNTIME_GID, 0o700),
        ('/var/lib/truf/host-agent', 0, 0, 0o700),
        ('/var/lib/truf/host-agent/backups', 0, 0, 0o700),
        ('/var/lib/truf/host-agent/operations', 0, 0, 0o700),
        ('/var/lib/truf/host-agent/results', 0, RUNTIME_GID, 0o750),
        ('/run/truf-postgres', RUNTIME_UID, RUNTIME_GID, 0o700),
    )
    for path, uid, gid, mode in layouts:
        _details(Path(path), directory=True, uid=uid, gid=gid, mode=mode)
    for executable in (
        '/usr/bin/python3', '/usr/bin/docker', '/usr/bin/systemctl',
        '/usr/bin/systemd-analyze', '/usr/bin/systemd-tmpfiles',
        '/usr/sbin/groupadd',
    ):
        _secure_executable(executable)
    for unit in UNITS:
        _details(SYSTEMD / unit, directory=False, uid=0, gid=0, mode=0o644)
        _same_file(SYSTEMD / unit, DEPLOY / unit)
    _details(TMPFILES, directory=False, uid=0, gid=0, mode=0o644)
    _same_file(TMPFILES, DEPLOY / 'truf-host-agent.conf')
    for script in SCRIPTS:
        _details(INSTALL_ROOT / script, directory=False, uid=0, gid=0, mode=0o755)
        _same_file(INSTALL_ROOT / script, DEPLOY / script)
    profile = _deployment_profile()
    for compose_file in profile['compose_files']:
        _read_source(PROJECT / compose_file)
    try:
        docker_socket = os.stat('/run/docker.sock', follow_symlinks=False)
    except OSError:
        raise InstallError('socket') from None
    if (
        not stat.S_ISSOCK(docker_socket.st_mode)
        or docker_socket.st_uid != 0
        or stat.S_IMODE(docker_socket.st_mode) & 0o002
    ):
        raise InstallError('socket')
    if require_socket:
        _validate_agent_socket()
    _run(('/usr/bin/python3', '-I', '-B', '-c', 'import psycopg, yaml'))
    _run(('/usr/bin/docker', 'compose', 'version'))
    _run(('/usr/bin/systemd-analyze', 'verify', *(str(SYSTEMD / unit) for unit in UNITS)))
    _validate_compose_projection(
        _capture(_compose_config_command(profile)), profile,
    )


def install():
    if sys.platform != 'linux' or not hasattr(os, 'geteuid') or os.geteuid() != 0:
        raise InstallError('root')
    _ensure_runtime_group()
    for unit in UNITS:
        if _unit_active(unit):
            _run(('/usr/bin/systemctl', 'stop', unit))
    _ensure_install_root()
    for script in SCRIPTS:
        _write(INSTALL_ROOT / script, _read_source(DEPLOY / script), uid=0, gid=0, mode=0o755, replace=True)
    for unit in UNITS:
        _write(SYSTEMD / unit, _read_source(DEPLOY / unit), uid=0, gid=0, mode=0o644, replace=True)
    _write(TMPFILES, _read_source(DEPLOY / 'truf-host-agent.conf'), uid=0, gid=0, mode=0o644, replace=True)
    _run(('/usr/bin/systemd-tmpfiles', '--create', str(TMPFILES)))
    _write(
        ACTIVE / 'config.yaml', _read_source(PROJECT / 'app/config.linux.yaml'),
        uid=RUNTIME_UID, gid=RUNTIME_GID, mode=0o600, replace=False,
    )
    _write(
        ACTIVE / 'secrets.yaml', b'{}\n', uid=RUNTIME_UID, gid=RUNTIME_GID,
        mode=0o600, replace=False,
    )
    validate(require_socket=False)
    _run(('/usr/bin/systemctl', 'daemon-reload'))
    _run(('/usr/bin/systemctl', 'enable', 'truf-host-agent.socket'))
    _run(('/usr/bin/systemctl', 'restart', 'truf-host-agent.socket'))
    _validate_agent_socket()


def main():
    if len(sys.argv) != 2 or sys.argv[1] not in ('install', 'validate'):
        raise InstallError('arguments')
    install() if sys.argv[1] == 'install' else validate()
    return 0


if __name__ == '__main__':
    try:
        result = main()
    except Exception as exc:
        print(
            'host-agent deployment failed (' + type(exc).__name__ + '); details withheld',
            file=sys.stderr,
        )
        result = 1
    raise SystemExit(result)