Check if screensaver is running before taking action

bin/kill_screensaver_graphic_program.sh

@@ -4,17 +4,19 @@ function help () {
    echo "Usage: $0 [OPTIONS] <file>"
    echo "Cleans up xscreensaver graphic programs left behind after killing xscreensaver itself."
    echo " -u     specify a username for which to kill xscreensaver artifacts"
+   echo " -d     dummy; just check if screensaver graphic program is running and exit."
    echo " -h     display this output"
    exit 1
 }
 
-while getopts ':u:' opt ; do
+while getopts ':u:dh' opt ; do
    case "$opt" in
       u)
          username="${OPTARG}"
          userid=`id -u "$username" 2>/dev/null` || { echo "$username is not a user on this system" && exit 1 ;}
          [ "$userid" -lt 1000 ] && echo "User $username seems to be a systemuser (uid: $userid). Please specify a normal user." && exit 1
          ;; 
+      d) dummy='true';;
       h) help ;;
       :)
          echo "$0: Must supply an argument to -$OPTARG." >&2
@@ -31,11 +33,13 @@ username="${username:-$USER}"
 homedir=`eval echo "~$username"`
 selected_graphic_nr=`grep selected "$homedir/.xscreensaver" | cut -f 2`
 regex='/\\n\\/'
-graphic_process_name=`cat "$homedir/.xscreensaver"  | awk "$regex && ++n == "$selected_graphic_nr" {getline; print; exit}" | cut -f 5 | cut -d ' ' -f 1 `
-echo "graphic_process_name: $graphic_process_name"
-[ -z "$graphic_process_name" ] && { echo Could not find selected graphic application && exit 1; }
-graphic_processes=`ps -U "$username" | grep "$graphic_process_name" | awk '{$1=$1};1' | cut -d ' ' -f 1 | tr '\n' ' ' ` || { echo "No screensaver graphic processes were found for $username" && exit 0; }
-echo "graphic_processes: $graphic_processes"
+graphic_program_name=`cat "$homedir/.xscreensaver"  | awk "$regex && ++n == "$selected_graphic_nr" {getline; print; exit}" | cut -f 5 | cut -d ' ' -f 1 `
+[ -z "$graphic_program_name" ] && { echo Could not find selected graphic application && exit 1; }
+echo "graphic_program: $graphic_program_name"
+
+graphic_processes=`ps -U "$username" | grep "$graphic_program_name" | awk '{$1=$1};1' | cut -d ' ' -f 1 | tr '\n' ' ' ` || { echo "No screensaver graphic processes were found for $username" && exit 0; }
 [ -z "$graphic_processes" ] && { echo "No xscreensaver graphic processes seem to be running for $username" && exit 0; }
+echo "graphic_processes: $graphic_processes"
+[ "$dummy" = 'true' ] && exit 0
 process_count=`echo $graphic_processes | wc -w`
 kill $graphic_processes && echo "killed $process_count screensaver graphic processes for $username"

bin/xscreensaver_yubilock.py

@@ -2,6 +2,7 @@
 
 #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
@@ -37,10 +38,10 @@ def setup_logger(logfile):
 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)
+                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
+            stderr=subprocess.PIPE) #Carefull with variables: shellcode injection
     out, err = process.communicate()
     if out != b'' :
         logger.info(out.decode('utf-8').rstrip())
@@ -48,6 +49,12 @@ def execute(command: str, shell_on: bool = False, background: bool = False):
         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)
+    if int(graphic_program_instances) > 0:
+        return True
+    return False
+
 def lock_screen():
     if args.dummy :
         return
@@ -59,9 +66,10 @@ def lock_screen():
 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)
+            shell_on=True)
 
     if xscreensaver_pid != 'null':
         execute('kill %s' % xscreensaver_pid, shell_on=True)
@@ -83,22 +91,30 @@ def get_yubikey_serials() -> int:
     except yubico.yubikey.YubiKeyError:
         pass
     except USBError as error:
-        logger.error('get_yubikey_serials() threw: %s', str(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():
+            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
 
-        if any(serial in yubikey_serials for serial in get_yubikey_serials()):
-            logger.debug('screen will be unlocked')
-            unlock_screen()
+        update_lock_state()
         return datetime.now()
     return last_yubikey_time
 
@@ -108,22 +124,22 @@ def device_removed(device: pyudev.Device, last_yubikey_time):
         # This is not nice, yubikey call can cause key to reconnect.
         # This cooldown prevents endless reconnect loops
 
-        if not any(serial in yubikey_serials for serial in get_yubikey_serials()):
-            logger.debug('screen will be locked')
-            lock_screen()
+        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.")
+            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')
+            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'
 
@@ -153,15 +169,13 @@ if __name__ == "__main__":
 
     execute('id > /tmp/id.txt &', shell_on=True, background = True)
     setup_logger(config.get("HOSTCONFIG", "logfile",
-             fallback="log.log"))
+        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
-    if not any(serial in yubikey_serials for serial in get_yubikey_serials()):
-        logger.debug('screen will be locked')
-        lock_screen()
+    update_lock_state()
 
     daemon(get_hid_event_monitor())