#!/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"])

screensaver = config.get("HOSTCONFIG", "screensaver", fallback="xscreensaver"))

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():
    if screensaver == 'xscreensaver':
        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
        if execute(f"DISPLAY=:0 xscreensaver-command -time | grep 'screen locked' >/dev/null 2>&1", shell_on=True): return True
    elif screensaver == 'gnome-screensaver':
        screensaver_state = execute("gnome-screensaver-command -q | tr -d '\n'", shell_on=True)
        if screensaver_state == 'The screensaver is active': return True
        if screensaver_state == 'The screensaver is inactive': return False
    return False

def lock_screen():
    if args.dummy :
        return
    if config.getboolean('USERCONFIG', 'remove_sudo_timestamp_when_locking', fallback=True):
        execute('sudo -K', shell_on=True)

    if screensaver == 'xscreensaver':
        execute('DISPLAY=:0 xscreensaver-command -lock', shell_on=True)
    elif screensaver == 'gnome-screensaver':
        execute('gnome-screensaver-command -l', shell_on=True)
    return

def unlock_screen():
    if args.dummy :
        return

    if screensaver == 'xscreensaver':
        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)
    elif screensaver == 'gnome-screensaver':
        execute(f"{script_dir}/loginctl-unlock.sh", 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():
            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"))


    if screensaver == 'xscreensaver':
        # 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())

    if screensaver == 'xscreensaver':
        # restart xscreensaver process before leaving
        execute('DISPLAY=:0 xscreensaver -no-splash&', shell_on=True, background = True)

    sleep(0.2)

    sys.exit(0)