Add handle

yubi_lock/api/user.py

@@ -14,3 +14,10 @@ def store_key(ident):
 
     with open(os.path.join(BASE_DIR, "data/keys.json"), 'w') as file:
         json.dump(existing_data, file)
+
+
+def has_ident(ident):
+    with open(os.path.join(BASE_DIR, "data/keys.json")) as file:
+        existing_data = json.load(file)
+
+    return ident in existing_data

yubi_lock/handle/cli.py

@@ -1,7 +1,22 @@
-import click
+import click, subprocess, os
+from yubi_lock.api.yubikey import get_all_yubikeys, generate_ident
+from yubi_lock.api.user import has_ident, get_username
 
 
 @click.command('handle', short_help='Handler for inserting and removing drives')
-@click.argument('type', type=click.Choice(["enable", "disable"]))
-def cli(type):
-    pass
+@click.argument('action', type=click.Choice(["enable", "disable"]))
+def cli(action):
+    ident = generate_ident(get_all_yubikeys(False)[0])
+    if not has_ident(ident):
+        return 1
+
+    session_id = subprocess.check_output(
+        "/bin/loginctl list-sessions | grep USER | awk '{print $1}'".replace('USER', get_username()),
+        shell=True
+    ).decode().replace('\n', '')
+    if action == 'enable':
+        os.system("/bin/loginctl lock-session {}".format(session_id))
+    elif action == 'disable':
+        os.system("/bin/loginctl unlock-session {}".format(session_id))
+    else:
+        return 1

yubi_lock/install/cli.py

@@ -5,8 +5,9 @@ from yubi_lock.api import BASE_DIR
 
 @click.command('install', short_help='Install required files for insertion detection')
 def cli():
-    exit_code = os.system("sudo cp {} {}".format(BASE_DIR + "../data/85-yubikey-screen-lock.rules", "/etc/udev/rules.d/85-yubikey-screen-lock.rules"))
+    exit_code = os.system("sudo cp {} {}".format(BASE_DIR + "/data/85-yubikey-screen-lock.rules", "/etc/udev/rules.d/85-yubikey-screen-lock.rules"))
     if exit_code == 0:
-        print('Export complete!')
+        print('Export complete! Refreshing...')
+        os.system("sudo udevadm control --reload-rules")
     else:
         print('Something went wrong!')