From 2d3fb50a4e50b6f8f378168d6c00256d73eca08c Mon Sep 17 00:00:00 2001
From: Scott Wallace <scott@wallace.sh>
Date: Mon, 13 Apr 2020 10:51:37 +0100
Subject: [PATCH] Refactor to run within a single Python script

---
 README.md         |  35 +++----
 bw_add_sshkeys.py | 261 ++++++++++++++++++++++++++++++++++++++++++++++
 ssh.py            |  68 ------------
 start.sh          |  12 ---
 4 files changed, 277 insertions(+), 99 deletions(-)
 create mode 100755 bw_add_sshkeys.py
 delete mode 100755 ssh.py
 delete mode 100755 start.sh

diff --git a/README.md b/README.md
index 24ba3b8..93cfc7f 100644
--- a/README.md
+++ b/README.md
@@ -1,31 +1,28 @@
 # Bitwarden SSH Agent
 
 ## Requirements
+* You need to have the [Bitwarden CLI tool](https://github.com/bitwarden/cli) installed and available in the `$PATH` as `bw`.
+* `ssh-agent` must be running in the current session.
 
-* You need to have the bitwarden cli `bw` installed
-* ssh-agent must be running in the current session
-
-## What it does?
-
-* connects to bitwarden using the bitwarden cli
-* looks for a folder called `ssh-agent`
-* loads the ssh key for each item in that folder
+## What does it do?
+Fetches SSH keys stored in Bitwarden vault and adds them to `ssh-agent`.
 
 ##  How to use it
-
-`$ ./start.sh`
-
-Fill in you login information 
+1. Run,
+   ```shell
+   ./bw_add_sshkeys.py
+   ```
+2. Enter your Bitwarden credentials, if a Bitwarden vault session is not already set.
+3. (optional) Enter your SSH keys' passphrases.
 
 
 ## Storing the keys in BitWarden
-
-1. Create a folder called 'ssh-agent'
-2. Add an new secure note to that folder
-3. Upload the private_key as an attachment
-4. add the custom field `private`, containing the private key filename
+1. Create a folder called `ssh-agent` (can be overridden on the command line).
+2. Add an new secure note to that folder.
+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 creating a file for it
+* 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
new file mode 100755
index 0000000..3f0ce34
--- /dev/null
+++ b/bw_add_sshkeys.py
@@ -0,0 +1,261 @@
+#!/usr/bin/env python3
+"""
+Extracts SSH keys from Bitwarden vault
+"""
+
+import argparse
+import getpass
+import json
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+
+
+def get_session():
+    """
+    Function to return a valid Bitwarden session
+    """
+    # Check for an existing, user-supplied Bitwarden session
+    try:
+        if os.environ['BW_SESSION']:
+            logging.debug('Existing Bitwarden session found')
+            return os.environ['BW_SESSION']
+    except KeyError:
+        pass
+
+    # Check if we're already logged in
+    proc = subprocess.Popen(
+        [
+            'bw',
+            'login',
+            '--check',
+            '--quiet'
+        ]
+    )
+    proc.wait()
+
+    if proc.returncode:
+        logging.debug('Not logged into Bitwarden')
+        operation = 'login'
+        credentials = [bytes(input('Bitwarden user: '), encoding='ascii')]
+    else:
+        logging.debug('Bitwarden vault is locked')
+        operation = 'unlock'
+        credentials = []
+
+    # Ask for the password
+    credentials.append(bytes(getpass.getpass('Bitwarden Vault password: '), encoding='ascii'))
+
+    proc = subprocess.Popen(
+        [
+            'bw',
+            '--raw',
+            '--nointeraction',
+            operation
+        ] + credentials,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    (stdout, stderr) = proc.communicate()
+
+    if proc.returncode:
+        logging.error(stderr.decode('utf-8'))
+        return None
+
+    return stdout.decode('utf-8')
+
+
+def get_folders(session, foldername):
+    """
+    Function to return the ID  of the folder that matches the provided name
+    """
+    logging.debug('Folder name: %s', foldername)
+
+    proc = subprocess.Popen(
+        [
+            'bw',
+            '--nointeraction',
+            'list',
+            'folders',
+            '--search', foldername,
+            '--session', session
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    (stdout, stderr) = proc.communicate()
+
+    if proc.returncode:
+        logging.error(stderr.decode('utf-8'))
+        return None
+
+    folders = json.loads(stdout)
+
+    if not folders:
+        logging.error('"%s" folder not found', foldername)
+        return None
+
+    # Do we have any folders
+    if len(folders) != 1:
+        logging.error('%d folders with the name "%s" found', len(folders), foldername)
+        return None
+
+    return folders[0]['id']
+
+
+def folder_items(session, folder_id):
+    """
+    Function to return items from a folder
+    """
+    logging.debug('Folder ID: %s', folder_id)
+
+    proc = subprocess.Popen(
+        [
+            'bw',
+            '--nointeraction',
+            'list',
+            'items',
+            '--folderid', folder_id,
+            '--session', session
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    (stdout, stderr) = proc.communicate()
+
+    if proc.returncode:
+        logging.error(stderr.decode('utf-8'))
+        return None
+
+    return json.loads(stdout)
+
+
+def add_ssh_keys(session, items, keyname):
+    """
+    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]
+        except IndexError:
+            logging.warning('No "%s" field found for item %s', keyname, 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]
+        except IndexError:
+            logging.warning(
+                'No attachment called "%s" found for item %s',
+                private_key_file,
+                item['name']
+            )
+            continue
+        logging.debug('Private key ID found')
+
+        if not ssh_add(session, item['id'], private_key_id):
+            logging.warning('Could not add key to the SSD agent')
+
+
+def ssh_add(session, item_id, key_id):
+    """
+    Function to get the key contents from the Bitwarden vault
+    """
+    logging.debug('Item ID: %s', item_id)
+    logging.debug('Key ID: %s', key_id)
+
+    # TODO: avoid temporary files, if possible
+    with tempfile.NamedTemporaryFile() as tmpfile:
+        proc = subprocess.Popen(
+            [
+                'bw',
+                '--nointeraction',
+                '--quiet',
+                'get',
+                'attachment', key_id,
+                '--itemid', item_id,
+                '--output', tmpfile.name,
+                '--session', session
+            ],
+            stderr=subprocess.PIPE
+        )
+        (_, stderr) = proc.communicate()
+        if proc.returncode:
+            logging.error(stderr.decode('utf-8'))
+            return False
+
+        logging.debug("Running ssh-add")
+
+        # CAVEAT: `ssh-add` provides no useful output, even with maximum verbosity
+        proc = subprocess.Popen(['ssh-add', tmpfile.name])
+        proc.wait()
+
+        if proc.returncode:
+            return False
+
+        return True
+
+
+if __name__ == '__main__':
+    def parse_args():
+        """
+        Function to parse command line arguments
+        """
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            '-d', '--debug',
+            action='store_true',
+            help='show debug output',
+        )
+        parser.add_argument(
+            '-f', '--foldername',
+            default='ssh-agent',
+            help='folder name to use to search for SSH keys',
+        )
+        parser.add_argument(
+            '-c', '--customfield',
+            default='private',
+            help='custom field name where private key filename is stored',
+        )
+
+        return parser.parse_args()
+
+
+    def main():
+        """
+        Main program logic
+        """
+
+        args = parse_args()
+
+        if args.debug:
+            loglevel = logging.DEBUG
+        else:
+            loglevel = logging.INFO
+
+        logging.basicConfig(level=loglevel)
+
+        logging.info('Getting Bitwarden session')
+        session = get_session()
+        if not session:
+            sys.exit(1)
+        logging.debug('Session = %s', session)
+
+        logging.info('Getting folder list')
+        folder_id = get_folders(session, args.foldername)
+        if not folder_id:
+            sys.exit(2)
+
+        logging.info('Getting folder items')
+        items = folder_items(session, folder_id)
+        if not items:
+            sys.exit(3)
+
+        logging.info('Attempting to add keys to ssh-agent')
+        add_ssh_keys(session, items, args.customfield)
+
+    main()
diff --git a/ssh.py b/ssh.py
deleted file mode 100755
index a68b429..0000000
--- a/ssh.py
+++ /dev/null
@@ -1,68 +0,0 @@
-import subprocess
-import os
-import sys
-import json
-import tempfile
-from urllib.parse import urlparse, urlencode
-from urllib.request import urlopen, Request
-from urllib.error import HTTPError
-
-try:
-    subprocess.check_output(['bw', 'logout'])
-except:
-    pass
-
-try:
-    session = subprocess.check_output(['bw', '--raw', 'login'])
-    session = ['--session', session]
-except:
-    print('Couldnt login!')
-    sys.exit(1)
-
-try:
-    folders = subprocess.check_output(['bw','list', 'folders', '--search', 'ssh-agent'] + session)
-    folders = json.loads(folders)
-    if not folders:
-        raise AttributeError
-    if len(folders) != 1:
-        raise ValueError
-except AttributeError:
-    print('Couldnt find ssh-agent folder!')
-    sys.exit(1)
-except ValueError:
-    print('More than one ssh-agent folder found!')
-    sys.exit(1)
-except:
-    print('Error retrieving folders.')
-    sys.exit(1)
-
-folder = folders[0]['id']
-
-try:
-    items = subprocess.check_output(['bw', 'list', 'items', '--folderid', folder, 'ssh-agent'] + session)
-    items = json.loads(items)
-except Exception as e:
-    print('Cant fint items.')
-    print(e)
-    sys.exit(1)
-
-keys = []
-directory = tempfile.TemporaryDirectory()
-temporary_file_name = os.path.join(directory.name + './private_key')
-
-try:
-    for item in items:
-        private_key_file = [k['value'] for k in item['fields'] if k['name'] == 'private' and k['type'] == 0][0]
-        
-        private_key_id = [k['id'] for k in item['attachments'] if k['fileName'] == private_key_file][0]
-
-        # would be nice if there was an option to retrieve the attachment file directly to the stdout
-        subprocess.check_output(['bw', 'get', 'attachment', private_key_id, '--itemid', item['id'], '--output', temporary_file_name] + session)
-        private_key = open(temporary_file_name, 'r').read()
-        os.remove(temporary_file_name)
-        keys.append({'private_key': private_key})
-except:
-    print('Something happened.')
-    sys.exit(1)
-
-print(';'.join([k['private_key'] for k in keys]))
diff --git a/start.sh b/start.sh
deleted file mode 100755
index 92cbdb7..0000000
--- a/start.sh
+++ /dev/null
@@ -1,12 +0,0 @@
-#!/bin/bash
-(
-KEYS=$(/usr/bin/env python ssh.py)
-IFS=';'
-read -d '' -ra SPLITKEYS < <(printf '%s;\0' "$KEYS")
-
-for i in ${SPLITKEYS[@]}
-do
-    ssh-add - <<< "${i}"
-done
-
-)
\ No newline at end of file