#!/usr/bin/python3 -I
"""Restore the existing GNOME Polkit agent if the shell/its native agent stops."""
import fcntl
import os
from pathlib import Path
import signal
import subprocess
import sys
import time

GNOME = '/usr/lib/polkit-gnome/polkit-gnome-authentication-agent-1'


def process_token(pid):
    try:
        path = Path('/proc') / str(pid)
        if path.stat().st_uid != os.getuid():
            return ''
        return (path / 'stat').read_text().rsplit(')', 1)[1].split()[19]
    except (OSError, IndexError):
        return ''


def agents():
    for entry in Path('/proc').iterdir():
        if not entry.name.isdecimal():
            continue
        try:
            if entry.stat().st_uid == os.getuid() and (entry / 'cmdline').read_bytes().split(b'\0')[0].decode() == GNOME:
                yield int(entry.name)
        except (OSError, UnicodeError):
            continue


def restore():
    if not list(agents()):
        subprocess.Popen([GNOME], start_new_session=True, stdin=subprocess.DEVNULL,
                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def main():
    if len(sys.argv) != 3 or sys.argv[1] not in ('acquire', 'heartbeat', 'watch', 'restore'):
        raise SystemExit(2)
    mode, pid = sys.argv[1], int(sys.argv[2])
    runtime = Path('/run/user') / str(os.getuid()) / 'windoesnot-auth'
    runtime.mkdir(mode=0o700, exist_ok=True)
    if runtime.is_symlink() or runtime.stat().st_uid != os.getuid():
        raise SystemExit('Invalid runtime directory')
    token = process_token(pid)
    if not token:
        raise SystemExit(1)
    lease = runtime / (str(pid) + '-' + token)
    owner = runtime / 'owner'
    if mode == 'acquire':
        owner.write_text(lease.name)
        lease.touch(mode=0o600)
        # Only replace the user's known fallback, never kill arbitrary agents.
        for agent in agents():
            os.kill(agent, signal.SIGTERM)
        for attempt in range(30):
            if not list(agents()):
                break
            time.sleep(0.05)
        subprocess.Popen([sys.executable, '-I', __file__, 'watch', str(pid)],
            start_new_session=True, stdin=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    elif mode == 'heartbeat':
        # A missing lease means fallback has already taken over.
        if not lease.exists():
            raise SystemExit(1)
        lease.touch()
    elif mode == 'restore':
        lease.unlink(missing_ok=True)
        restore()
    else:
        with (runtime / (lease.name + '.lock')).open('w') as lock:
            try:
                fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError:
                return
            while process_token(pid) == token:
                try:
                    if time.time() - lease.stat().st_mtime > 15:
                        break
                except FileNotFoundError:
                    break
                time.sleep(1)
            lease.unlink(missing_ok=True)
            if not owner.exists() or owner.read_text() != lease.name:
                return
            # Native agent may still be alive after an event-loop hang. GNOME
            # refuses duplicate registration; retry until the session recovers.
            restore()


if __name__ == '__main__':
    main()
