diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..da87b55 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.pyenv/ +.mypy_cache/ diff --git a/README.md b/README.md index 93cfc7f..87a3577 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,3 @@ Fetches SSH keys stored in Bitwarden vault and adds them to `ssh-agent`. 3. Upload the private key as an attachment. 4. Add the custom field `private` (can be overridden on the command line), containing the file name of the private key attachment. 5. Repeat steps 2-4 for each subsequent key - - -## Improvements to be made -* Find a way to extract the attachment from bitwarden in memory, instead of using a temporary file. diff --git a/bw_add_sshkeys.py b/bw_add_sshkeys.py index efcb0f4..6acc636 100755 --- a/bw_add_sshkeys.py +++ b/bw_add_sshkeys.py @@ -8,17 +8,18 @@ import json import logging import os import subprocess +from typing import Any, Callable, Dict, List from pkg_resources import parse_version -def memoize(func): +def memoize(func: Callable[..., Any]) -> Callable[..., Any]: """ Decorator function to cache the results of another function call """ - cache = dict() + cache: Dict[Any, Callable[..., Any]] = {} - def memoized_func(*args): + def memoized_func(*args: Any) -> Any: if args in cache: return cache[args] result = func(*args) @@ -29,7 +30,7 @@ def memoize(func): @memoize -def bwcli_version(): +def bwcli_version() -> str: """ Function to return the version of the Bitwarden CLI """ @@ -43,7 +44,7 @@ def bwcli_version(): @memoize -def cli_supports(feature): +def cli_supports(feature: str) -> bool: """ Function to return whether the current Bitwarden CLI supports a particular feature @@ -55,18 +56,18 @@ def cli_supports(feature): return False -def get_session(): +def get_session() -> str: """ Function to return a valid Bitwarden session """ # Check for an existing, user-supplied Bitwarden session - session = os.environ.get('BW_SESSION') - if session is not None: + session = os.environ.get('BW_SESSION', '') + if session: logging.debug('Existing Bitwarden session found') return session # Check if we're already logged in - proc_logged = subprocess.run(['bw', 'login', '--check', '--quiet']) + proc_logged = subprocess.run(['bw', 'login', '--check', '--quiet'], check=True) if proc_logged.returncode: logging.debug('Not logged into Bitwarden') @@ -81,10 +82,15 @@ def get_session(): universal_newlines=True, check=True, ) - return proc_session.stdout + session = proc_session.stdout + logging.info( + 'To re-use this BitWarden session run: export BW_SESSION="%s"', + session, + ) + return session -def get_folders(session, foldername): +def get_folders(session: str, foldername: str) -> str: """ Function to return the ID of the folder that matches the provided name """ @@ -101,55 +107,66 @@ def get_folders(session, foldername): if not folders: logging.error('"%s" folder not found', foldername) - return None + return '' # Do we have any folders if len(folders) != 1: logging.error('%d folders with the name "%s" found', len(folders), foldername) - return None + return '' - return folders[0]['id'] + return str(folders[0]['id']) -def folder_items(session, folder_id): +def folder_items(session: str, folder_id: str) -> List[Dict[str, Any]]: """ Function to return items from a folder """ logging.debug('Folder ID: %s', folder_id) proc_items = subprocess.run( - [ 'bw', 'list', 'items', '--folderid', folder_id, '--session', session], + ['bw', 'list', 'items', '--folderid', folder_id, '--session', session], stdout=subprocess.PIPE, universal_newlines=True, check=True, ) - return json.loads(proc_items.stdout) + + data: List[Dict[str, Any]] = json.loads(proc_items.stdout) + + return data -def add_ssh_keys(session, items, keyname): +def add_ssh_keys(session: str, items: List[Dict[str, Any]], keyname: str) -> None: """ Function to attempt to get keys from a vault item """ for item in items: try: - private_key_file = [k['value'] for k in item['fields'] - if k['name'] == keyname and k['type'] == 0][0] + private_key_file = [ + k['value'] + for k in item['fields'] + if k['name'] == keyname and k['type'] == 0 + ][0] except IndexError: logging.warning('No "%s" field found for item %s', keyname, item['name']) continue - except KeyError as e: - logging.debug('No key "%s" found in item %s - skipping', e.args[0], item['name']) + except KeyError as error: + logging.debug( + 'No key "%s" found in item %s - skipping', error.args[0], item['name'] + ) continue logging.debug('Private key file declared') try: - private_key_id = [k['id'] for k in item['attachments'] - if k['fileName'] == private_key_file][0] + private_key_id = [ + k['id'] + for k in item['attachments'] + if k['fileName'] == private_key_file + ][0] except IndexError: logging.warning( 'No attachment called "%s" found for item %s', private_key_file, - item['name'] + item['name'], ) continue logging.debug('Private key ID found') @@ -160,20 +177,24 @@ def add_ssh_keys(session, items, keyname): logging.warning('Could not add key to the SSH agent') -def ssh_add(session, item_id, key_id): +def ssh_add(session: str, item_id: str, key_id: str) -> None: """ Function to get the key contents from the Bitwarden vault """ logging.debug('Item ID: %s', item_id) logging.debug('Key ID: %s', key_id) - proc_attachment = subprocess.run([ + proc_attachment = subprocess.run( + [ 'bw', 'get', - 'attachment', key_id, - '--itemid', item_id, + 'attachment', + key_id, + '--itemid', + item_id, '--raw', - '--session', session + '--session', + session, ], stdout=subprocess.PIPE, universal_newlines=True, @@ -195,31 +216,34 @@ def ssh_add(session, item_id, key_id): if __name__ == '__main__': - def parse_args(): + + def parse_args() -> argparse.Namespace: """ Function to parse command line arguments """ parser = argparse.ArgumentParser() parser.add_argument( - '-d', '--debug', + '-d', + '--debug', action='store_true', help='show debug output', ) parser.add_argument( - '-f', '--foldername', + '-f', + '--foldername', default='ssh-agent', help='folder name to use to search for SSH keys', ) parser.add_argument( - '-c', '--customfield', + '-c', + '--customfield', default='private', help='custom field name where private key filename is stored', ) return parser.parse_args() - - def main(): + def main() -> None: """ Main program logic """ @@ -246,9 +270,9 @@ if __name__ == '__main__': logging.info('Attempting to add keys to ssh-agent') add_ssh_keys(session, items, args.customfield) - except subprocess.CalledProcessError as e: - if e.stderr: - logging.error('`%s` error: %s', e.cmd[0], e.stderr) - logging.debug('Error running %s', e.cmd) + except subprocess.CalledProcessError as error: + if error.stderr: + logging.error('`%s` error: %s', error.cmd[0], error.stderr) + logging.debug('Error running %s', error.cmd) main()