From b714bddd94fb0aed75e35e31642fe5107a015769 Mon Sep 17 00:00:00 2001
From: Jake Howard <git@theorangeone.net>
Date: Thu, 19 May 2016 22:24:55 +0100
Subject: [PATCH] Add handle

---
 yubi_lock/api/user.py    |  7 +++++++
 yubi_lock/handle/cli.py  | 23 +++++++++++++++++++----
 yubi_lock/install/cli.py |  5 +++--
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/yubi_lock/api/user.py b/yubi_lock/api/user.py
index ade507e..1ac7c94 100644
--- a/yubi_lock/api/user.py
+++ b/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
diff --git a/yubi_lock/handle/cli.py b/yubi_lock/handle/cli.py
index 76b10d7..e5800af 100644
--- a/yubi_lock/handle/cli.py
+++ b/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
diff --git a/yubi_lock/install/cli.py b/yubi_lock/install/cli.py
index 7da7ee2..8de759b 100644
--- a/yubi_lock/install/cli.py
+++ b/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!')