#!/usr/bin/python3 -I
"""Fixed, polkit-authorized Luxwarp session service setup entry point."""

import hashlib
import importlib.util
import os
from pathlib import Path
import pwd
import re
import stat
import subprocess
import sys

SELF = Path('/usr/libexec/luxwarp/session-service-setup')
INSTALLER = Path('/usr/share/luxwarp/session-service/install-session-service.py')
BUNDLE = INSTALLER.parent
LOGINCTL = '/usr/bin/loginctl'
SYSTEMCTL = '/usr/bin/systemctl'
UNIT = 'luxwarp-session.service'
ARTIFACTS = ('update_safety.py', 'manifest.json', 'lw-host', 'luxwarp-session-supervisor.py',
             'luxwarp-session.service')
SAFE_ENV = {'PATH': '/usr/bin:/bin', 'LC_ALL': 'C'}


def trusted_file(path):
    if not path.is_absolute():
        raise ValueError('installed path is not absolute')
    current = Path('/')
    for part in path.parts[1:-1]:
        current /= part
        info = current.lstat()
        if not stat.S_ISDIR(info.st_mode) or info.st_uid != 0 or info.st_mode & 0o022:
            raise ValueError(f'untrusted installed directory: {current}')
    info = path.lstat()
    if (not stat.S_ISREG(info.st_mode) or info.st_nlink != 1
            or info.st_uid != 0 or info.st_mode & 0o022):
        raise ValueError(f'untrusted installed file: {path}')
    return info


def caller():
    raw = os.environ.get('PKEXEC_UID', '')
    if not raw.isdecimal() or raw.startswith('0') or len(raw) > 10:
        raise ValueError('pkexec did not identify an ordinary caller')
    uid = int(raw)
    if uid < 1000 or uid == 65534:
        raise ValueError('session setup requires an ordinary desktop user')
    account = pwd.getpwuid(uid)
    if (not re.fullmatch(r'[a-z_][a-z0-9_-]{0,31}', account.pw_name)
            or account.pw_dir != f'/home/{account.pw_name}'
            or account.pw_uid != uid):
        raise ValueError('session setup requires a local account under /home')
    return account


def properties(kind, ident, names):
    command = [LOGINCTL, f'show-{kind}', ident, '--no-pager']
    for name in names:
        command.extend(('-p', name))
    result = subprocess.run(command, check=True, capture_output=True, timeout=5,
                            cwd='/', env=SAFE_ENV)
    values = {}
    for line in result.stdout.decode('utf-8', 'strict').splitlines():
        key, sep, value = line.partition('=')
        if sep and key in names and key not in values:
            values[key] = value
    if set(values) != set(names):
        raise ValueError('incomplete seat information')
    return values


def current_x11_seat0(account):
    active = properties('seat', 'seat0', ('ActiveSession',))['ActiveSession'].split()
    ident = active[0] if active else ''
    if not re.fullmatch(r'[A-Za-z0-9_-]{1,64}', ident):
        raise ValueError('no active seat0 session')
    state = properties('session', ident,
                       ('Active', 'State', 'Class', 'Type', 'Remote', 'Name', 'User', 'Seat'))
    user = state['User'].split()[0] if state['User'].split() else ''
    expected = {'Active': 'yes', 'State': 'active', 'Class': 'user', 'Type': 'x11',
                'Remote': 'no', 'Name': account.pw_name, 'User': str(account.pw_uid),
                'Seat': 'seat0'}
    if {**state, 'User': user} != expected:
        raise ValueError('caller does not own the local active X11 seat0 session')


def load_installer():
    trusted_file(SELF)
    if os.path.abspath(__file__) != str(SELF):
        raise ValueError('setup must run from its installed path')
    for path in (INSTALLER, *(BUNDLE / name for name in ARTIFACTS)):
        trusted_file(path)
    spec = importlib.util.spec_from_file_location('luxwarp_installed_session_installer', INSTALLER)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def main():
    if len(sys.argv) != 2 or sys.argv[1] not in ('--enable', '--disable'):
        raise ValueError('expected exactly --enable or --disable')
    if sys.platform != 'linux' or os.geteuid() != 0:
        raise ValueError('session setup requires pkexec on Linux')
    account = caller()
    installer = load_installer()
    owner = installer.installed_owner()
    if owner is not None and owner != account.pw_name:
        raise ValueError('existing Luxwarp service belongs to another user')
    if sys.argv[1] == '--disable':
        if owner is None:
            raise ValueError('no installed Luxwarp session service for this user')
        subprocess.run([SYSTEMCTL, 'stop', UNIT], check=True, timeout=20,
                       stdin=subprocess.DEVNULL, env=SAFE_ENV, cwd='/')
        subprocess.run([SYSTEMCTL, 'disable', UNIT], check=True, timeout=20,
                       stdin=subprocess.DEVNULL, env=SAFE_ENV, cwd='/')
        print('Luxwarp sign-in remote access disabled')
        return
    current_x11_seat0(account)
    manifest = BUNDLE / 'manifest.json'
    digest = hashlib.sha256(manifest.read_bytes()).hexdigest()
    installer.load_bundle(BUNDLE, digest)
    subprocess.run(['/usr/bin/python3', '-I', str(INSTALLER), '--bundle', str(BUNDLE),
                    '--user', account.pw_name, '--manifest-sha256', digest,
                    '--replace-running-host'], check=True,
                   stdin=subprocess.DEVNULL, env=SAFE_ENV, cwd='/')


if __name__ == '__main__':
    try:
        main()
    except (OSError, ValueError, KeyError, UnicodeError, subprocess.SubprocessError) as error:
        print(f'Luxwarp session setup failed: {error}', file=sys.stderr)
        sys.exit(1)
