#!/usr/bin/env python3

#TODO:
# - resolve exploit where key is plugged in and out quickly, resulting in an unlocked state

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

config = ConfigParser()
config.read("config.ini")

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 lock_screen():
    if args.dummy :
        return
    if config.getboolean('HOSTCONFIG', '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)

    # 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.error('get_yubikey_serials() threw: %s', str(error))

    if not serials:
        logger.info('no yubikey connected')

    return serials

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()
        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

        if not any(serial in yubikey_serials for serial in get_yubikey_serials()):
            logger.debug('screen will be locked')
            lock_screen()
        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
    if not any(serial in yubikey_serials for serial in get_yubikey_serials()):
        logger.debug('screen will be locked')
        lock_screen()

    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)