#!/usr/bin/env python3
"""Deny desktop-user power actions on Ubuntu/Kylin; run as root in the guest.

Writes account-scoped polkit and D-Bus rules. Never sends a valid power request.
Root/PVE maintenance remains available. Query uses login1 Can* methods only.
"""
import argparse
import json
import os
from pathlib import Path
import pwd
import re
import shutil
import subprocess
import time
from xml.sax.saxutils import quoteattr


def run(command):
    return subprocess.run(command, capture_output=True, text=True, timeout=20,
                          env={**os.environ, 'LC_ALL': 'C', 'SYSTEMD_COLORS': '0'})


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['apply', 'query', 'restore'])
    parser.add_argument('--user', required=True)
    parser.add_argument('--backup', default='/var/lib/vdi/power-policy')
    args = parser.parse_args()
    if os.geteuid() != 0:
        parser.error('Run as root in the guest, not on the terminal or PVE node')
    if not re.fullmatch(r'[a-z_][a-z0-9_-]*', args.user):
        parser.error('Invalid desktop user')
    account = pwd.getpwnam(args.user)
    if account.pw_uid == 0:
        parser.error('Refusing to restrict root')
    backend = 'pkla' if run(['pkaction', '--version']).stdout.strip().endswith('0.105') else 'rules'
    if backend == 'pkla':
        target = Path('/etc/polkit-1/localauthority/90-mandatory.d/99-vdi-power.pkla')
        content = ('[VDI desktop power restriction]\nIdentity=unix-user:' + args.user + '\n'
                   'Action=org.freedesktop.login1.power-off*;org.freedesktop.login1.reboot*;'
                   'org.freedesktop.login1.halt*;org.freedesktop.login1.suspend*;'
                   'org.freedesktop.login1.hibernate*;org.freedesktop.login1.hybrid-sleep*;'
                   'org.freedesktop.consolekit.system.stop*;org.freedesktop.consolekit.system.restart*;'
                   'org.freedesktop.upower.suspend;org.freedesktop.upower.hibernate\n'
                   'ResultAny=no\nResultInactive=no\nResultActive=no\n')
    else:
        target = Path('/etc/polkit-1/rules.d/00-vdi-power.rules')
        content = '''polkit.addRule(function(action, subject) {
    if (subject.user === %s &&
        /^org\\.freedesktop\\.login1\\.(power-off|reboot|halt|suspend|hibernate|hybrid-sleep)(-|$)/.test(action.id)) {
        return polkit.Result.NO;
    }
});
''' % json.dumps(args.user)
    backup = Path(args.backup)
    manifest = backup / 'manifest.json'
    dbus_path = Path('/etc/dbus-1/system.d/zz-vdi-desktop-power.conf')
    methods = ('PowerOff', 'Reboot', 'Halt', 'Suspend', 'Hibernate', 'HybridSleep',
               'SuspendThenHibernate', 'ScheduleShutdown')
    capability_methods = ('CanPowerOff', 'CanReboot', 'CanHalt', 'CanSuspend',
                          'CanHibernate', 'CanHybridSleep', 'CanSuspendThenHibernate')
    members = [m for name in methods for m in (name, name + 'WithFlags')] + list(capability_methods)
    dbus_content = ('<!DOCTYPE busconfig PUBLIC "-//freedesktop//DTD D-BUS Bus Configuration 1.0//EN"\n'
                    ' "http://www.freedesktop.org/standards/dbus/1.0/busconfig.dtd">\n'
                    '<busconfig>\n  <policy user=' + quoteattr(args.user) + '>\n' + ''.join(
                        '    <deny send_destination="org.freedesktop.login1" '
                        'send_interface="org.freedesktop.login1.Manager" send_member="' + m + '"/>\n'
                        for m in members) + '  </policy>\n</busconfig>\n')
    dbus_manifest = backup / 'dbus-manifest.json'
    if args.action == 'apply':
        backup.mkdir(parents=True, exist_ok=True, mode=0o700)
        if not manifest.exists():
            exists = target.exists()
            if exists:
                shutil.copy2(target, backup / 'original')
            manifest.write_text(json.dumps({'user': args.user, 'path': str(target), 'existed': exists}))
        saved = json.loads(manifest.read_text())
        if saved['user'] != args.user or saved['path'] != str(target):
            raise RuntimeError('Backup belongs to a different target')
        target.parent.mkdir(parents=True, exist_ok=True)
        temporary = target.with_suffix(target.suffix + '.tmp')
        temporary.write_text(content)
        temporary.chmod(0o644)
        temporary.replace(target)
        # Active-session authorization can still grant Reboot. Enforce the
        # boundary at the bus, including the menu's capability queries.
        if not dbus_manifest.exists():
            if dbus_path.exists():
                shutil.copy2(dbus_path, backup / 'dbus-original')
            dbus_manifest.write_text(json.dumps({'user': args.user, 'existed': dbus_path.exists()}))
        if json.loads(dbus_manifest.read_text())['user'] != args.user:
            raise RuntimeError('D-Bus backup belongs to a different user')
        dbus_tmp = dbus_path.with_suffix('.conf.tmp')
        dbus_tmp.write_text(dbus_content)
        dbus_tmp.chmod(0o644)
        dbus_tmp.replace(dbus_path)
        reload_result = run(['busctl', '--system', 'call', 'org.freedesktop.DBus',
                             '/org/freedesktop/DBus', 'org.freedesktop.DBus', 'ReloadConfig'])
        if reload_result.returncode:
            raise RuntimeError('D-Bus reload failed: ' + reload_result.stderr)
        time.sleep(2)  # Both polkit backends monitor their policy directories.
    elif args.action == 'restore':
        saved = json.loads(manifest.read_text())
        if saved['user'] != args.user or saved['path'] != str(target):
            raise RuntimeError('Backup belongs to a different target')
        dbus_saved = json.loads(dbus_manifest.read_text())
        if dbus_saved['user'] != args.user:
            raise RuntimeError('D-Bus backup belongs to another user')
        if saved['existed']:
            shutil.copy2(backup / 'original', target)
        elif target.exists():
            target.unlink()
        if dbus_saved['existed']:
            shutil.copy2(backup / 'dbus-original', dbus_path)
        elif dbus_path.exists():
            dbus_path.unlink()
        reload_result = run(['busctl', '--system', 'call', 'org.freedesktop.DBus',
                             '/org/freedesktop/DBus', 'org.freedesktop.DBus', 'ReloadConfig'])
        if reload_result.returncode:
            raise RuntimeError('D-Bus reload failed: ' + reload_result.stderr)
        time.sleep(2)
    capabilities = {}
    for method in capability_methods:
        result = run(['runuser', '-u', args.user, '--', 'busctl', '--system', 'call',
                      'org.freedesktop.login1', '/org/freedesktop/login1',
                      'org.freedesktop.login1.Manager', method])
        capabilities[method] = {'code': result.returncode, 'value': result.stdout.strip(),
                                'error': result.stderr.strip()}
    active = target.exists() and target.read_text() == content
    dbus_active = dbus_path.exists() and dbus_path.read_text() == dbus_content
    bus_denials = {}
    if args.action == 'apply':
        # Deliberately use a STRING argument. Every real power method here
        # requires a boolean, flags, or a different signature, so even if the
        # bus rule failed, logind rejects this before any power action runs.
        # AccessDenied (rather than InvalidArgs) proves the bus blocks delivery.
        for method in methods:
            result = run(['runuser', '-u', args.user, '--', 'busctl', '--system', 'call',
                          'org.freedesktop.login1', '/org/freedesktop/login1',
                          'org.freedesktop.login1.Manager', method, 's', 'VDI_INVALID_SIGNATURE'])
            bus_denials[method] = {'code': result.returncode, 'error': result.stderr.strip()}
    print(json.dumps({'action': args.action, 'user': args.user, 'backend': backend,
                      'path': str(target), 'backup': str(backup), 'policy_matches': active,
                      'dbus_policy_matches': dbus_active, 'dbus_denials': bus_denials,
                      'boot_id': Path('/proc/sys/kernel/random/boot_id').read_text().strip(),
                      'capabilities': capabilities}))
    if args.action == 'apply' and (not active or not dbus_active or any(
            not (entry['code'] and 'Access denied' in entry['error'])
            for entry in capabilities.values()) or any(
            not entry['code'] or not any(message in entry['error'] for message in
                                        ('Rejected send message', 'Access denied'))
            for entry in bus_denials.values())):
        raise SystemExit('Policy installed but denial verification failed; inspect capabilities')


if __name__ == '__main__':
    main()