yubicheck.py

import subprocess
import sys
import ctypes
import os
import json
import msvcrt
from time import sleep

DEFAULT_CHECK_INTERVAL_SECONDS = 10

def get_base_dir():
    """Return the directory where config should be read from.

    When packaged to an exe, use the executable folder.
    Otherwise, use the script folder.
    """
    if getattr(sys, "frozen", False):
        return os.path.dirname(sys.executable)
    return os.path.dirname(os.path.abspath(__file__))

CONFIG_PATH = os.path.join(get_base_dir(), "yubicheck_config.json")

def load_config(config_path):
    """Load check interval and whitelisted serials from JSON config."""
    check_interval_seconds = DEFAULT_CHECK_INTERVAL_SECONDS
    whitelisted_serials = set()

    try:
        with open(config_path, "r", encoding="utf-8") as f:
            cfg = json.load(f)
    except FileNotFoundError:
        print("WARN: config file not found: " + config_path)
        return check_interval_seconds, whitelisted_serials
    except json.JSONDecodeError:
        print("WARN: invalid JSON config: " + config_path)
        return check_interval_seconds, whitelisted_serials
    except OSError as ex:
        print("WARN: could not read config: " + str(ex))
        return check_interval_seconds, whitelisted_serials

    interval_value = cfg.get("check_interval_seconds", DEFAULT_CHECK_INTERVAL_SECONDS)
    try:
        interval_value = int(interval_value)
        if interval_value > 0:
            check_interval_seconds = interval_value
        else:
            print("WARN: check_interval_seconds must be > 0, using default")
    except (TypeError, ValueError):
        print("WARN: check_interval_seconds is not a valid integer, using default")

    serials_value = cfg.get("whitelisted_serials", [])
    if isinstance(serials_value, list):
        whitelisted_serials = {str(s).strip() for s in serials_value if str(s).strip()}
    else:
        print("WARN: whitelisted_serials must be a JSON array")

    return check_interval_seconds, whitelisted_serials

def save_config(config_path, check_interval_seconds, whitelisted_serials):
    """Persist config to JSON so updates survive exe builds."""
    cfg = {
        "check_interval_seconds": int(check_interval_seconds),
        "whitelisted_serials": sorted(whitelisted_serials),
    }

    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=2)
        f.write("\n")

def prompt_add_to_whitelist(new_serials, timeout_seconds):
    """Return True if user confirms adding detected serials within timeout."""
    if not sys.stdin or not sys.stdin.isatty():
        print("Interactive prompt skipped (no console input available)")
        return False

    print("Unwhitelisted YubiKey(s) detected on startup: " + ", ".join(sorted(new_serials)))
    print(
        "Press 'y' within "
        + str(timeout_seconds)
        + " seconds to add to whitelist, or 'n' to ignore."
    )

    steps = timeout_seconds * 10
    for _ in range(steps):
        if msvcrt.kbhit():
            key = msvcrt.getwch().strip().lower()
            if key == "y":
                print("User approved whitelist update")
                return True
            if key == "n":
                print("User rejected whitelist update")
                return False
        sleep(0.1)

    print("No user input received before timeout")
    return False

CHECK_INTERVAL_SECONDS, WHITELISTED_SERIALS = load_config(CONFIG_PATH)

first=True
startupinfo = subprocess.STARTUPINFO()
startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW

print("Using config: " + CONFIG_PATH)
print("Check interval seconds: " + str(CHECK_INTERVAL_SECONDS))
print("Whitelisted serial count: " + str(len(WHITELISTED_SERIALS)))

def get_connected_serials():
    """Return a set of serial numbers reported by ykman."""
    p = subprocess.Popen(
        ["ykman", "list", "--serials"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        startupinfo=startupinfo,
        text=True,
    )
    stdout, _ = p.communicate()

    if p.returncode != 0:
        return None

    serials = set()
    for line in stdout.splitlines():
        serial = line.strip()
        if serial:
            serials.add(serial)
    return serials

while(1):
    connected_serials = get_connected_serials()
    if connected_serials is None:
        print("ERROR: ykman query failed")
        yubi_found = False
    else:
        print("Connected serials: " + ", ".join(sorted(connected_serials)) if connected_serials else "Connected serials: (none)")
        if not WHITELISTED_SERIALS:
            # Fail-safe: if no whitelist is configured, do not treat any key as authorized.
            yubi_found = False
        else:
            yubi_found = any(serial in WHITELISTED_SERIALS for serial in connected_serials)

        if first and (not yubi_found) and connected_serials:
            new_serials = connected_serials - WHITELISTED_SERIALS
            if new_serials and prompt_add_to_whitelist(new_serials, 10):
                WHITELISTED_SERIALS.update(new_serials)
                try:
                    save_config(CONFIG_PATH, CHECK_INTERVAL_SECONDS, WHITELISTED_SERIALS)
                    print("Whitelist updated in config")
                except OSError as ex:
                    print("WARN: could not save config: " + str(ex))

                yubi_found = any(serial in WHITELISTED_SERIALS for serial in connected_serials)

    if yubi_found:
        print("Whitelisted YubiKey found")
    else:
        if(first):
            sys.exit(0)
        print("locking")
        ctypes.windll.user32.LockWorkStation()
    first=False
    sleep(CHECK_INTERVAL_SECONDS)