#!/usr/bin/env python3
# systemd-lock-handler -- proxy between systemd-logind's "Lock" signal and your
# favourite screen lock command

import argparse
import dbus
import dbus.mainloop.glib
from gi.repository import GLib
import os
import sys

def setup_signal(session_id, signal_handler):
    bus = dbus.SystemBus()
    manager = bus.get_object("org.freedesktop.login1", "/org/freedesktop/login1")
    # yecch
    manager = dbus.Interface(manager, "org.freedesktop.login1.Manager")
    session_path = manager.GetSession(session_id)
    session = bus.get_object("org.freedesktop.login1", session_path)
    session.connect_to_signal("Lock", signal_handler)
    return session

def handler_dbus_fdo():
    print("Locking session using D-Bus")
    bus = dbus.SessionBus()
    screensaver = bus.get_object("org.freedesktop.ScreenSaver", "/ScreenSaver")
    screensaver.Lock()

def handler_external(lock_command):
    print("Locking session using %r" % lock_command[0])
    os.spawnvp(os.P_NOWAIT, lock_command[0], lock_command)

parser = argparse.ArgumentParser()
parser.add_argument("--dbus", action="store_true",
                    help="use the 'ScreenSaver' D-Bus API for locking")
parser.add_argument("command", nargs="*",
                    help="run the specified command to lock the screen")
args = parser.parse_args()

if not (args.dbus or args.command):
    exit("error: Either D-Bus mode or an executable must be specified.")
elif (args.dbus and args.command):
    # We *could* very easily allow both, but usually it's pointless.
    exit("error: D-Bus mode cannot be used together with command mode.")

try:
    session_id = os.environ["XDG_SESSION_ID"]
except KeyError:
    # TODO: Figure out what to do when we're running out of systemd --user.
    # How does xss-lock handle that? (Edit: It just gives up.)
    exit("error: $XDG_SESSION_ID not set; are you using pam_systemd?")

dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)

if args.dbus:
    print("Using ScreenSaver D-Bus API")
    listener = setup_signal(session_id,
                            handler_dbus_fdo)
elif args.command:
    print("Using external command %r" % args.command[0])
    listener = setup_signal(session_id,
                            lambda: handler_external(args.command))

print("Waiting for lock signals on session %s" % session_id)

try:
    loop = GLib.MainLoop()
    loop.run()
except KeyboardInterrupt:
    exit()
