Merge pull request joaojacome#11 from Yamakaky/master

Improve subprocess handling

bw_add_sshkeys.py

@@ -4,13 +4,10 @@
 """
 
 import argparse
-import getpass
 import json
 import logging
 import os
 import subprocess
-import sys
-import tempfile
 
 from pkg_resources import parse_version
 
@@ -36,20 +33,13 @@ def bwcli_version():
     """
     Function to return the version of the Bitwarden CLI
     """
-    proc = subprocess.Popen(
-        [
-            'bw',
-            '--version'
-        ],
-        stdout=subprocess.PIPE
+    proc_version = subprocess.run(
+        ['bw', '--version'],
+        stdout=subprocess.PIPE,
+        text=True,
+        check=True,
     )
-
-    (stdout, _) = proc.communicate()
-
-    if proc.returncode:
-        raise RuntimeError('Unable to fetch Bitwarden CLI version')
-
-    return stdout.decode('utf-8')
+    return proc_version.stdout
 
 
 @memoize
@@ -70,53 +60,28 @@ 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
+    session = os.environ.get('BW_SESSION')
+    if session is not None:
+        logging.debug('Existing Bitwarden session found')
+        return session
 
     # Check if we're already logged in
-    proc = subprocess.Popen(
-        [
-            'bw',
-            'login',
-            '--check',
-            '--quiet'
-        ]
-    )
-    proc.wait()
+    proc_logged = subprocess.run(['bw', 'login', '--check', '--quiet'])
 
-    if proc.returncode:
+    if proc_logged.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(
-        list(filter(None, [
-            'bw',
-            '--raw',
-            (None, '--nointeraction')[cli_supports('nointeraction')],
-            operation
-        ] + credentials)),
+    proc_session = subprocess.run(
+        ['bw', '--raw', operation],
         stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
+        text=True,
+        check=True,
     )
-    (stdout, stderr) = proc.communicate()
-
-    if proc.returncode:
-        logging.error(stderr.decode('utf-8'))
-        return None
-
-    return stdout.decode('utf-8')
+    return proc_session.stdout
 
 
 def get_folders(session, foldername):
@@ -125,25 +90,14 @@ def get_folders(session, foldername):
     """
     logging.debug('Folder name: %s', foldername)
 
-    proc = subprocess.Popen(
-        list(filter(None, [
-            'bw',
-            (None, '--nointeraction')[cli_supports('nointeraction')],
-            'list',
-            'folders',
-            '--search', foldername,
-            '--session', session
-        ])),
+    proc_folders = subprocess.run(
+        ['bw', 'list', 'folders', '--search', foldername, '--session', session],
         stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
+        text=True,
+        check=True,
     )
-    (stdout, stderr) = proc.communicate()
 
-    if proc.returncode:
-        logging.error(stderr.decode('utf-8'))
-        return None
-
-    folders = json.loads(stdout)
+    folders = json.loads(proc_folders.stdout)
 
     if not folders:
         logging.error('"%s" folder not found', foldername)
@@ -163,25 +117,13 @@ def folder_items(session, folder_id):
     """
     logging.debug('Folder ID: %s', folder_id)
 
-    proc = subprocess.Popen(
-        list(filter(None, [
-            'bw',
-            (None, '--nointeraction')[cli_supports('nointeraction')],
-            'list',
-            'items',
-            '--folderid', folder_id,
-            '--session', session
-        ])),
+    proc_items = subprocess.run(
+        [ 'bw', 'list', 'items', '--folderid', folder_id, '--session', session],
         stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
+        text=True,
+        check=True,
     )
-    (stdout, stderr) = proc.communicate()
-
-    if proc.returncode:
-        logging.error(stderr.decode('utf-8'))
-        return None
-
-    return json.loads(stdout)
+    return json.loads(proc_items.stdout)
 
 
 def add_ssh_keys(session, items, keyname):
@@ -209,7 +151,9 @@ def add_ssh_keys(session, items, keyname):
             continue
         logging.debug('Private key ID found')
 
-        if not ssh_add(session, item['id'], private_key_id):
+        try:
+            ssh_add(session, item['id'], private_key_id)
+        except subprocess.SubprocessError:
             logging.warning('Could not add key to the SSH agent')
 
 
@@ -220,36 +164,31 @@ def ssh_add(session, item_id, key_id):
     logging.debug('Item ID: %s', item_id)
     logging.debug('Key ID: %s', key_id)
 
-    # FIXME: avoid temporary files, if possible (StringIO ?)
-    with tempfile.NamedTemporaryFile() as tmpfile:
-        proc = subprocess.Popen(
-            list(filter(None, [
-                'bw',
-                (None, '--nointeraction')[cli_supports('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
+    proc_attachment = subprocess.run([
+            'bw',
+            'get',
+            'attachment', key_id,
+            '--itemid', item_id,
+            '--raw',
+            '--session', session
+        ],
+        stdout=subprocess.PIPE,
+        text=True,
+        check=True,
+    )
+    ssh_key = proc_attachment.stdout
+
+    logging.debug("Running ssh-add")
+
+    # CAVEAT: `ssh-add` provides no useful output, even with maximum verbosity
+    subprocess.run(
+        ['ssh-add', '-'],
+        input=ssh_key,
+        # Works even if ssh-askpass is not installed
+        env=dict(os.environ, SSH_ASKPASS_REQUIRE="never"),
+        text=True,
+        check=True,
+    )
 
 
 if __name__ == '__main__':
@@ -291,23 +230,22 @@ def main():
 
         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)
+        try:
+            logging.info('Getting Bitwarden session')
+            session = get_session()
+            logging.debug('Session = %s', session)
+
+            logging.info('Getting folder list')
+            folder_id = get_folders(session, args.foldername)
+
+            logging.info('Getting folder items')
+            items = folder_items(session, folder_id)
+
+            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)
 
     main()