Bitwarden script to log in & retrieve a password

bwlogin.py

#!/usr/bin/env python3

"""
INTRO:
This is a Python script to log in & retrieve your SSH keys from Bitwarden. It will save your session token to a
local file and invalidate that file after a certain timeframe. This script will NOT start your SSH agent for you,
but it will expect that your SSH agent is already running and properly configured. If in doubt, run the following
command in your shell before this script to start and configure an agent first: `source <(ssh-agent)`

USAGE:
1) Save the passphrase of your key to Bitwarden in a "note" object, then use the Bitwarden CLI to find their IDs.
2) For the sake of being pretty, this code uses blessings. It should be the only pip package you need to install - 
    go ahead and do that.
3) Fill in the variables at the top of this file in caps (TOKENFILE, USER, KEYS) appropriately.
4) Profit!
"""

import time
import json
import os
import subprocess
import shlex
from blessings import Terminal
from getpass import getpass
import sys

# Path where the Bitwarden session token should be created and kept by this script.
TOKENFILE = f'{os.path.join(os.path.dirname(os.path.realpath(__file__)), "session.tkn")}'

# Your Bitwarden login email
USER = '[email protected]'

# A list of keys you want to load in the following format:
# (path_to_private_key, id_in_bitwarden)
KEYS: list[tuple] = [
    ('/home/you/.ssh/id_rsa_example_1', 'aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'),
    ('/home/you/.ssh/id_rsa_example_2', 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb', )]

# Time in seconds that a session token should be considered valid
TIMEFRAME = 24 * 60 * 60 # one day

""" Should not need any modifications past this line """    

term = Terminal()

def main():
    try:
        token = get_token()

        print()
        for i, key in enumerate(KEYS):
            cr()
            print(f'Loading key {i+1}/{len(KEYS)}...')
            passphrase = retrieve_pw(key[1], token)
            subprocess.run(f'''export SSH_ASKPASS=/bin/cat ; unset DISPLAY ; echo '{passphrase}' | setsid -w ssh-add '{key[0]}' >/dev/null 2>/dev/null''', stdout=subprocess.PIPE, shell=True)
        cr()
        print(f'Loading key {len(KEYS)}/{len(KEYS)}... {term.green}{term.underline}Done.')

    except KeyboardInterrupt:
        sys.exit(130)

        
def cr():
    sys.stdout.write('\033[F\033[K')


def get_token() -> str:
    # If the token is less than a day old and valid, we can just use that.
    if os.path.isfile(TOKENFILE) and not time.time() - os.path.getmtime(TOKENFILE) > TIMEFRAME:
        token = open(TOKENFILE, 'r').readline()
        response = call_shell(f'bw status --session {token}')
        try:
            if json.loads(response)['status'] == 'unlocked':
                return token
            raise None
        except:
            print(f'{term.yellow("!")} Session token invalid. Re-generating...')

    # Otherwise, log out to be safe with BW before logging in again
    call_shell('bw logout')
    if os.path.isfile(TOKENFILE):
        os.remove(TOKENFILE)

    token = ''
    while token == '':
        pw = getpass(f'{term.green("?")} {term.bold("Master password:")} {term.italic("[input is hidden]")} ')
        cr()
        print(f'{term.green("?")} {term.bold("Master password:")} {term.italic("[hidden]")}')

        p = subprocess.Popen(shlex.split(
            f'bw login {USER}'), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        p.stdin.write(bytes(pw.encode('ascii') + b'\n'))
        p.stdin.flush()

        while p.returncode == None:
            p.poll()
        
        if b'incorrect' in p.stderr.read():
            cr()
            print(f'''{term.red('✘')} {term.bold('Master password:')} {term.italic(f'[incorrect]')}''')
            continue

        p.stdout.flush()
        for line in p.stdout.readlines():
            line = str(line)
            if line.find('export BW_SESSION') != -1:
                start = line.find('"')
                end = line.rfind('"') + 1
                token = line[start:end]
                break
        else:
            print()
            exit()

    # Save that bad boy to the token file
    with open(TOKENFILE, 'w') as fout:
        fout.write(token)

    # Keep the terminal pretty
    cr()
    print(f'''{term.green('✔')} {term.bold('Master password:')} {term.italic(f'[correct]')}''')

    return token


def retrieve_pw(key: str, token: str) -> str:
    """Return a password for a specified key

    Args:
        id (str): Bitwarden Secure Note object ID (find with `bw list items --search 'nameofitem')
        token (str): Valid Bitwarden session token

    Returns:
        str: Note contents for `id`
    """
    return json.loads(call_shell(f'bw get item --session {token} {key}'))['notes']


def call_shell(cmd: str) -> str:
    """Calls shell within child process and returns stdout

    Args:
        cmd (str): Command to run

    Returns:
        str: Stdout of `cmd`
    """
    res = b''
    with subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE) as p:
        for line in p.stdout:
            res += line

    return str(res, 'ascii')


if __name__ == '__main__':
    main()