Dotfiles_Generator/templates/rofi/rofi-bitwarden.py

367 lines
11 KiB
Python
Executable File

{# symlink{~/bin/rofi-bitwarden.py} #}
#!/usr/bin/env python3
# TODO: totp
import os
import sys
import json
import shlex
import subprocess
from io import StringIO
from dataclasses import dataclass
from urllib.parse import urlparse
from functools import cache, partial
from typing import Callable, Tuple
# Time until the vault locks itself
# Resets upon every retrieval
# Set to 0 to disble autolocking
BW_SESSION_AUTOLOCK_TIMEOUT = 3600
# Change the key under which the secret is store using keyctl
BW_SESSION_SECRET_NAME = 'bw-session'
'''
Utils
'''
stderr = partial(print, file=sys.stderr)
class Clipboard():
@staticmethod
def copy(value):
run_shell_command('xclip -in -selection clipboard', input_data=value, disable_stdout=True, disable_stderr=True)
@staticmethod
def paste() -> str:
return run_shell_command('xclip -out -selection clipboard').stdout
'''
Shell Commands
'''
@dataclass
class ShellOutput:
status_code: int
stdout: str
stderr: str
def run_shell_command(cmd, input_data=None, raise_on_error=False, disable_stdout=False, disable_stderr=False):
process = subprocess.run(
cmd if type(cmd) == list else shlex.split(cmd),
text=True,
input=input_data,
stdout=None if disable_stdout else subprocess.PIPE,
stderr=None if disable_stderr else subprocess.PIPE,
check=raise_on_error,
)
return ShellOutput(
process.returncode,
None if disable_stdout else process.stdout.strip(),
None if disable_stderr else process.stderr.strip(),
)
'''
Keyctl
'''
class KeyctlSecret():
def __init__(self, key, timeout = 0, keyring = '@u'):
self.key = key
self.timeout = timeout
self.keyring = keyring
def get_key_id(self):
output = run_shell_command(['keyctl', 'request', 'user', self.key])
if output.status_code > 0:
return None
self.reset_timeout()
return output.stdout
def store(self, value):
run_shell_command(['keyctl', 'add', 'user', self.key, value, self.keyring], raise_on_error=True)
self.reset_timeout()
def reset_timeout(self):
if self.timeout > 0:
run_shell_command(['keyctl', 'timeout', self.key, str(self.timeout)])
def get_value(self):
key_id = self.get_key_id()
if key_id is None:
return None
return run_shell_command(['keyctl', 'pipe', key_id], raise_on_error=True).stdout
'''
Bitwarden
'''
class Bitwarden():
def __init__(self):
self.session_secret = KeyctlSecret(BW_SESSION_SECRET_NAME, BW_SESSION_AUTOLOCK_TIMEOUT)
def ensure_vault_unlocked(self):
self._get_session_token()
def get_all_entries(self):
return self._get_bw_entries()
def get_all_entries_by_uri(self, uri):
return self._get_bw_entries('--url', uri)
def _get_session_token(self):
# Check if the secret is already set
if self.session_secret.get_key_id() is None:
# If not, ask for password through rofi
password = show_rofi_pwd_prompt()
# Use the password to unlock BW vault
session_token = self._unlock(password)
# Store the new session token
self.session_secret.store(session_token)
# Return the value stored by the secret manager
return self.session_secret.get_value()
def _unlock(self, password):
bw_unlock_output = run_shell_command(['bw', 'unlock', password, '--raw'])
if bw_unlock_output.status_code > 0:
stderr(f'Bitwarden error: {bw_unlock_output.stderr}')
exit(1)
return bw_unlock_output.stdout
@cache
def _get_bw_entries(self, *filter_args):
session_token = self._get_session_token()
output = run_shell_command(['bw', '--session', session_token, 'list', 'items'] + list(filter_args))
if output.status_code > 0:
stderr(output.stderr)
exit(1)
return json.loads(output.stdout)
'''
Menu States
'''
@dataclass
class MenuEntry():
key: str
label: str
class MainMenu():
def __init__(self, app, bw):
self.app = app
self.bw = bw
self.desired_domain = None
def get_title(self):
return 'Bitwarden'
def get_entries(self):
if run_shell_command('xdotool getwindowfocus getwindowclassname').stdout == 'qutebrowser':
old_contents = Clipboard.paste()
run_shell_command('qutebrowser :yank')
self.desired_domain = Clipboard.paste()
Clipboard.copy(old_contents)
entries = []
if self.desired_domain is not None:
entries += [MenuEntry(
'search-domain',
f'Search on URL: {self.desired_domain}',
)]
entries += [
MenuEntry('list-name', 'List by Name'),
MenuEntry('list-url', 'List by URL'),
]
return entries
def handle_input(self, key, typed_input):
if key == 'list-name':
self.app.change_menu_state(ListByNameMenu(self.app, self.bw))
elif key == 'list-url':
self.app.change_menu_state(ListByURLMenu(self.app, self.bw))
elif key == 'search-domain':
all_entries = self.bw.get_all_entries()
filtered_entries = self.bw.get_all_entries_by_uri(self.desired_domain)
filtered_ids = [entry['id'] for entry in filtered_entries]
indices = [i for i, entry in enumerate(all_entries) if entry['id'] in filtered_ids]
self.app.change_menu_state(ViewListMatchMenu(self.app, self.bw, indices))
else:
stderr("Unreachable input")
class ViewEntryMenu:
def __init__(self, app, bw, entry_index):
self.app = app
self.bw = bw
self.entry = self.bw.get_all_entries()[int(entry_index)]
self.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'],
}
def get_title(self):
return f'Bitwarden - Copy {self.entry["name"]}'
def get_entries(self):
entries = []
for key in self.selection_options:
try:
if self.selection_options[key](self.entry) is not None:
entries += [MenuEntry(
key,
key,
)]
except:
pass
return entries
def handle_input(self, key, typed_input):
value = self.selection_options[key](self.entry)
Clipboard.copy(value)
self.app.set_finished()
class ViewListMatchMenu:
def __init__(self, app, bw, possible_indices):
self.app = app
self.bw = bw
self.possible_indices = possible_indices
def get_title(self):
return 'Bitwarden - List Matches'
def get_entries(self):
raw_entries = self.bw.get_all_entries()
menu_entries = []
for index in self.possible_indices:
entry = raw_entries[index]
key = f'{entry["name"]} ({entry["login"]["username"]})'
menu_entries += [MenuEntry(
str(index),
key,
)]
return menu_entries
def handle_input(self, key, typed_input):
self.app.change_menu_state(ViewEntryMenu(self.app, self.bw, int(key)))
class ListByAnyFieldMenu():
def __init__(self, app, bw, entry_to_keys_func):
self.app = app
self.bw = bw
self.entry_to_keys_func = entry_to_keys_func
self.key_to_entry_indices_map = {}
def get_title(self):
return 'Bitwarden - List'
def get_entries(self):
raw_entries = self.bw.get_all_entries()
if len(self.key_to_entry_indices_map.keys()) == 0:
for index, entry in enumerate(raw_entries):
if 'login' not in entry.keys():
continue
for key in self.entry_to_keys_func(entry):
if key is None:
continue
if key in self.key_to_entry_indices_map:
self.key_to_entry_indices_map[key] += [index]
else:
self.key_to_entry_indices_map[key] = [index]
menu_entries = []
for index, key in enumerate(sorted(self.key_to_entry_indices_map.keys())):
menu_entries += [MenuEntry(
str(index),
key,
)]
return menu_entries
def handle_input(self, key, typed_input):
possible_indices = self.key_to_entry_indices_map[sorted(self.key_to_entry_indices_map.keys())[int(key)]]
if len(possible_indices) > 1:
self.app.change_menu_state(ViewListMatchMenu(self.app, self.bw, possible_indices))
else:
self.app.change_menu_state(ViewEntryMenu(self.app, self.bw, possible_indices[0]))
class ListByNameMenu(ListByAnyFieldMenu):
def __init__(self, app, bw):
super().__init__(app, bw, lambda e: [e['name']])
class ListByURLMenu(ListByAnyFieldMenu):
def __init__(self, app, bw):
super().__init__(app, bw, lambda e: [urlparse(uri['uri']).hostname for uri in e['login']['uris']])
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=True):
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)
class App():
def __init__(self):
self.bw = Bitwarden()
self.menu_state = MainMenu(self, self.bw)
self._is_finished = False
def is_finished(self):
return self._is_finished
def set_finished(self):
self._is_finished = True
def render(self):
entries = self.menu_state.get_entries()
(index, extra_text) = show_rofi_entries([entry.label for entry in entries], title=self.menu_state.get_title())
self.menu_state.handle_input(entries[index].key, extra_text)
def change_menu_state(self, new_state):
self.menu_state = new_state
if __name__ == '__main__':
app = App()
while not app.is_finished():
app.render()