#!/usr/bin/env python3 #TODO: # - resolve exploit where key is plugged in and out quickly, resulting in an unlocked state # - Add async scheduler that runs update_lock_state() every 1 minute or so import os import sys import subprocess import argparse from datetime import datetime from configparser import ConfigParser from time import sleep import pyudev import yubico import logzero from logzero import logger from usb.core import USBError script_dir = os.path.dirname(os.path.realpath(__file__)) home_dir = os.path.expanduser("~") config = ConfigParser() config.read([f"{script_dir}/config.ini", f"{home_dir}/.yubilock"]) yubikey_serials = config["USERCONFIG"]["yubikey_serial"].split(',') # Convert stringlist to intlist for i, serial in enumerate(yubikey_serials): yubikey_serials[i] = int(serial) def setup_logger(logfile): logzero.loglevel(10 if args.verbose else 100) loglevel = config.getint("HOSTCONFIG", "loglevel", fallback=10) logzero.logfile(logfile, loglevel=loglevel, maxBytes=40960, backupCount=4) logger.info('script started') def execute(command: str, shell_on: bool = False, background: bool = False): if background : subprocess.run(command, shell=shell_on, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False) return process = subprocess.Popen(command, shell=shell_on, stdout=subprocess.PIPE, stderr=subprocess.PIPE) #Carefull with variables: shellcode injection out, err = process.communicate() if out != b'' : logger.info(out.decode('utf-8').rstrip()) if err != b'' : logger.error(err.decode('utf-8').rstrip()) return '' if out == b'' else out.decode('utf-8') def screensaver_running(): graphic_program_instances = execute(f"{script_dir}/kill_screensaver_graphic_program.sh -d | grep graphic_processes | wc -l", shell_on=True) return int(graphic_program_instances) > 0 def just_woke_up(): seconds_since_wakeup = execute(f"{script_dir}/seconds_since_wakeup.sh", shell_on=True) return int(seconds_since_wakeup) < 60 def lock_screen(): if args.dummy : return if config.getboolean('USERCONFIG', 'remove_sudo_timestamp_when_locking', fallback=True): execute('sudo -K', shell_on=True) execute('DISPLAY=:0 xscreensaver-command -lock', shell_on=True) return def unlock_screen(): if args.dummy : return logger.info('xscreen process to be killed:') xscreensaver_pid = execute(r'ps -A | grep -oPm 1 "\d{2,}(?=\s.+xscreensaver)" || echo null', shell_on=True) if xscreensaver_pid != 'null': execute('kill %s' % xscreensaver_pid, shell_on=True) execute(f"{script_dir}/kill_screensaver_graphic_program.sh", shell_on=True) # restart xscreensaver process execute('DISPLAY=:0 xscreensaver -no-splash&', shell_on=True, background = True) return def get_yubikey_serials() -> int: serials = [] try: skip = 0 while skip < 255: yubikey = yubico.find_yubikey(debug = False, skip = skip) serials.append(yubikey.serial()) logger.info('yubikey connected, serial= %s', yubikey.serial()) skip += 1 except yubico.yubikey.YubiKeyError: pass except USBError as error: logger.info('get_yubikey_serials() threw: %s', str(error)) if not serials: logger.info('no yubikey connected') return serials def update_lock_state(): if any(serial in yubikey_serials for serial in get_yubikey_serials()): if screensaver_running() or just_woke_up(): logger.debug('screen will be unlocked') unlock_screen() else: if not screensaver_running(): logger.debug('screen will be locked') lock_screen() def device_added(device, last_yubikey_time): logger.debug('device connected: %s', device.sys_name) if (datetime.now() - last_yubikey_time).total_seconds() > 0.3 : # This is not nice, yubikey call can cause key to reconnect. # This cooldown prevents endless reconnect loops update_lock_state() return datetime.now() return last_yubikey_time def device_removed(device: pyudev.Device, last_yubikey_time): logger.debug('device removed: %s', device.sys_name) if (datetime.now() - last_yubikey_time).total_seconds() > 0.3 : # This is not nice, yubikey call can cause key to reconnect. # This cooldown prevents endless reconnect loops update_lock_state() return datetime.now() return last_yubikey_time def get_args(): parser = argparse.ArgumentParser(description="Deamon that monitors for a specific YubiKey \ serial number. When the YubiKey is connected, any running xscreensaver instance will be \ unlocked. When it is removed, xscreensaver will lock again.") parser.add_argument('-d', '--dummy', action='store_true', \ help='don\'t actually lock and unlock the display') parser.add_argument('-v', '--verbose', action='store_true', help='increase output verbosity') return parser.parse_args() def daemon(monitor): last_yubikey_time = datetime.fromtimestamp(0) lastdevice = 'hidraw-2' try: for device in iter(monitor.poll, None): if device.sys_name[:6] == 'hidraw': if abs(int(device.sys_name[6:]) - int(lastdevice[6:])) == 1 : continue lastdevice = device.sys_name if device.action == 'add': last_yubikey_time = device_added(device, last_yubikey_time) elif device.action == 'remove': last_yubikey_time = device_removed(device, last_yubikey_time) else: logger.info('device ' + device.action + ': ' + device.sys_name) except KeyboardInterrupt: logger.info('script was terminiated by KeyboardInterrupt') def get_hid_event_monitor(): context = pyudev.Context() monitor = pyudev.Monitor.from_netlink(context) monitor.filter_by('hidraw') #human interface device connected/disconnected return monitor if __name__ == "__main__": args = get_args() setup_logger(config.get("HOSTCONFIG", "logfile", fallback="log.log")) # start xscreensaver process execute('DISPLAY=:0 xscreensaver -no-splash&', shell_on=True, background = True) # Lock the machine if no key is inserted when the script is started update_lock_state() daemon(get_hid_event_monitor()) # restart xscreensaver process before leaving execute('DISPLAY=:0 xscreensaver -no-splash&', shell_on=True, background = True) sleep(0.2) sys.exit(0)