Finish first version of bitwarden script

templates/rofi/rofi-bitwarden.py

@@ -4,65 +4,214 @@
 import os
 import sys
 import json
+import shlex
+import subprocess
 
-BW_SESSION_FILE = os.path.join('/var/run/user', str(os.getuid()), 'bitwarden-session')
+from io import StringIO
+from dataclasses import dataclass
+from urllib.parse import urlparse
+from functools import cache, partial
 
-def does_bw_session_exist():
+from typing import Callable, Tuple
-    return os.path.exists(BW_SESSION_FILE)
 
-def show_rofi_pwd_prompt():
+# Time until the vault locks itself
-    return os.popen('rofi -dmenu -password -i -no-fixed-num-lines -p "Master Password" -theme ~/.config/rofi/themes/askpass.rasi').read().strip()
+# Resets upon every retrieval
+# Set to 0 to disble autolocking
+BW_SESSION_AUTOLOCK_TIMEOUT = 0
 
-def get_bw_session():
+'''
-    if not does_bw_session_exist():
+Utils
-        return None
+'''
+stderr = partial(print, file=sys.stderr)
 
-    with open(BW_SESSION_FILE, 'r') as file:
+def copy_to_clipboard(value):
-        return file.read()
+    run_shell_command('xclip -in -selection clipboard', input_data=value)
 
-def is_logged_in():
+'''
-    return os.system('bw login --check > /dev/null') == 0
+Shell Commands
+'''
+@dataclass
+class ShellOutput:
+    status_code: int
+    stdout: str
+    stderr: str
 
-def login():
+def run_shell_command(cmd, input_data=None):
-    os.spawnl(os.P_WAIT, '/usr/bin/alacritty', 'alacritty', '-e', 'sh', '-c', 'bw login')
+    process = subprocess.run(
-
+        shlex.split(cmd),
-def unlock():
+        text=True,
-    password = show_rofi_pwd_prompt()
+        input=input_data,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
     
-    if len(password) > 0:
+    return ShellOutput(
-        # TODO: pipe in the password
+        process.returncode,
-        os.system(f'bw unlock "{password}" --raw > {BW_SESSION_FILE}')
+        process.stdout.strip(),
+        process.stderr.strip(),
+    )
 
-def lock():
+'''
-    if not does_bw_session_exist():
+Keyctl
-        return False
+'''
+KEYCTL_BW_SESSION_KEY='bw-session'
 
-    return os.system(f'bw --session "{get_bw_session()}" --nointeraction lock > /dev/null') == 0
+def get_bw_session_token():
+    request_output = run_shell_command(f'keyctl request user {KEYCTL_BW_SESSION_KEY}')
+    if request_output.status_code > 0:
+        # Ask for password
+        password = show_rofi_pwd_prompt()
+        bw_unlock_output = unlock(password)
 
-def is_vault_unlocked():
+        if bw_unlock_output.status_code > 0:
-    if not does_bw_session_exist():
+            stderr(f'Bitwarden error: {bw_unlock_output.stderr}')
-        return False
+            exit(1)
 
-    return os.system(f'bw --session "{get_bw_session()}" --nointeraction list items > /dev/null 2>&1') == 0
+        add_output = run_shell_command(f'keyctl add user {KEYCTL_BW_SESSION_KEY} "{bw_unlock_output.stdout}" @u')
+        if add_output.status_code > 0:
+            stderr('Could not store session token')
+            exit(1)
+        key_id = add_output.stdout
+    else:
+        key_id = request_output.stdout
 
-def debug(msg):
+    # Maybe invalidate on timeout
-    os.spawnl(os.P_WAIT, '/usr/bin/alacritty', 'alacritty', '-e', 'sh', '-c', f'echo "{msg}" && cat -')
+    if BW_SESSION_AUTOLOCK_TIMEOUT > 0:
+        print(run_shell_command(f'keyctl timeout {key_id} "{BW_SESSION_AUTOLOCK_TIMEOUT}"'))
 
+    pipe_output = run_shell_command(f'keyctl pipe {key_id}')
+
+    if pipe_output.status_code > 0:
+        stderr('Could not fetch session token')
+        exit(1)
+
+    return pipe_output.stdout
+
+'''
+Bitwarden
+'''
+def unlock(password):
+    return run_shell_command(f'bw unlock "{password}" --raw')
+
+@cache
+def get_all_entries():
+    session_token = get_bw_session_token()
+
+    output = run_shell_command(f'bw --session "{session_token}" list items')
+
+    if output.status_code > 0:
+        stderr(output.stderr)
+        exit(1)
+
+    return json.loads(output.stdout)
+
+def get_entries_by_filter(entry_to_keys):
+    """This function returns:
+    { key: str, entries: int[] }[]
+    """
+
+    output_entries = {}
+    entries = get_all_entries()
+
+    for entry_index, entry in enumerate(entries):
+        if 'login' not in entry.keys():
+            continue
+        
+        for key in entry_to_keys(entry):
+            if key is None:
+                continue
+
+            if key in output_entries:
+                output_entries[key] += [entry_index]
+            else:
+                output_entries[key] = [entry_index]
+
+    true_output = []
+    for key in output_entries:
+        true_output += [{
+            'key': key,
+            'entries': output_entries[key],
+        }]
+    return sorted(true_output, key=lambda e: e['key'])
+
+'''
+UI
+'''
+def show_rofi_pwd_prompt():
+    output = run_shell_command('rofi -dmenu -password -i -no-fixed-num-lines -p "Master Password" -theme ~/.config/rofi/themes/askpass.rasi')
+    if output.status_code > 0:
+        return None
+    return output.stdout
+
+def show_rofi_entries(entries, title="Bitwarden", exit_on_close=False):
+    output = run_shell_command(f'rofi -dmenu -i -p "{title}" -format "i:f"', input_data='\n'.join(entries))
+
+    if output.status_code > 0:
+        if exit_on_close:
+            exit(0)
+        return (None, None)
+
+    [selection_index, selection_text] = output.stdout.split(':')
+
+    return (int(selection_index), selection_text)
+
+'''
+Main
+'''
 def main():
-    if not is_logged_in():
+    get_bw_session_token()
-        login()
-        if not is_logged_in():
-            print('You did nog log in', file=sys.stderr)
-            exit(1)
 
-    if not is_vault_unlocked():
+    main_menu_entries = [
-        unlock()
+        # (key, display, function)
-        if not is_vault_unlocked():
+        (
-            print('You did not unlock the vault', file=sys.stderr)
+            'list-by-name',
-            exit(1)
+            'List by Name',
+            lambda e: [e['name']]
+        ), (
+            'list-by-uri',
+            'List by URI',
+            lambda e: [urlparse(uri['uri']).hostname for uri in e['login']['uris']]
+        )
+    ]
 
-    print("Signed in and unlocked, cool")
+    '''Main Menu'''
-    debug("Done")
+    (index, _) = show_rofi_entries([entry[1] for entry in main_menu_entries], exit_on_close=True)
+
+    '''List of items'''
+    selection = main_menu_entries[index]
+
+    entry_index_map = get_entries_by_filter(selection[2])
+    (index, _) = show_rofi_entries([entry['key'] for entry in entry_index_map], title=f'Bitwarden - {selection[1]}', exit_on_close=True)
+
+    '''List of items matching key'''
+    entries = get_all_entries()
+    selected_entries = entry_index_map[index]['entries']
+    
+    if len(selected_entries) == 1:
+        index = 0
+    else:
+        (index, _) = show_rofi_entries([f'{entries[index]["name"]} ({entries[index]["login"]["username"]})' for index in selected_entries], exit_on_close=True)
+
+    '''List copy options'''
+    selection_options = {
+        'Password': lambda e: e['login']['password'],
+        'TOTP': lambda e: e['login']['totp'],
+        'Username': lambda e: e['login']['username'],
+        'URI': lambda e: e['login']['uris'][0]['uri'],
+    }
+
+    selected_entry = entries[selected_entries[index]]
+    display_options = []
+    for key in selection_options:
+        try:
+            value = selection_options[key](selected_entry)
+            if value is not None:
+                display_options += [key]
+        except:
+            pass
+
+    (index, _) = show_rofi_entries(display_options, title=f'Bitwarden - Copy for {selected_entry["name"]}', exit_on_close=True)
+    copy_to_clipboard(selection_options[display_options[index]](selected_entry))
 
 if __name__ == '__main__':
     main()