7
\$\begingroup\$

While as a developer I naturally came across python many times in my life, I cannot recall writing any programs in python. So I'm sure my python is not very idiomatic. Let's fix it!

This is a script to automatically unseal Hashicorp Vault It is meant to run from within a docker container, but could be run stand-alone. It does not take any command line parameters, but takes input from either environment variables (a common method of passing configuration to a docker container) or from a config file, which I choose to be in json format.

The source code can be found here and also follows.

The three inputs are:

  • The url of vault you want to automatically unseal (e.g. https://vault.rocks)
  • The number of seconds to sleep between checks if the vault is sealed
  • A list of unseal keys (you probably want a single key, since if you do not care for manual unsealing, it's unlikely that you created more than one)

In addition, you also can specify an environment variable to enable debug mode which turns on more verbose output.

import os
import sys
import json
import time
import datetime
from pprint import pprint
import requests

version = "0.1"
envPrefix = "VU_"
config = "vault-unseal.json"

def PrintWithTimestamp(string):
  timestamp = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S')
  sys.stdout.write(f"[{timestamp}] {string}\n")

def PrintDebug(string):
  if "VU_DEBUG" in os.environ:
    PrintWithTimestamp(f"DEBUG: {string}")

def PrintParameterError(name):
  PrintWithTimestamp(f"Error: {name} is not specified. Specify it either in {config} file, or as enviroment variable prefixed with {envPrefix}")

def NormalizePrefix(name, string):
  return name + string[len(envPrefix+name):]

def ReadSingleSetting(name):
  result = None
  try:
    with open(config) as data_file:    
      data = json.load(data_file)
    result = data[name]
  except (FileNotFoundError, KeyError):
    pass
  if envPrefix+name in os.environ:
    result = os.environ[envPrefix+name]
  return result


def ReadMultiSetting(name):
  result = None
  try:
    with open(config) as data_file:    
      data = json.load(data_file)
    result = { key:data[key] for key in filter(lambda x: x.startswith(name), data.keys()) }
  except (FileNotFoundError):
    pass

  result = result if result else {}

  prefix = envPrefix + name
  envResult = { NormalizePrefix(name,key):os.environ[key] for key in filter(lambda x: x.lower().startswith(prefix.lower()), list(os.environ.keys())) }

  finalResult = list({**result,**envResult}.values())
  finalResult = finalResult if len(finalResult) else None
  return finalResult

def ReadSetting(name):
  if name.endswith("*"):
    return ReadMultiSetting(name[0:-1])
  else:
    return ReadSingleSetting(name)   

PrintWithTimestamp(f"vault-unseal.py version {version}")

if "VU_DEBUG" in os.environ:
  PrintDebug("Dumping environment block:")
  pprint(dict(os.environ))

addressUrl = ReadSetting("ADDRESS_URL")
timeIntervalSeconds = int(ReadSetting("TIME_INTERVAL_SECONDS"))
unsealKeys = ReadSetting("UNSEAL_KEY_*")

if not addressUrl:
  PrintParameterError("ADDRESS_URL")
  sys.exit(1)

if not timeIntervalSeconds:
  PrintParameterError("TIME_INTERVAL_SECONDS")
  sys.exit(1)

if not unsealKeys:
  PrintParameterError("UNSEAL_KEY_*")
  sys.exit(1)

PrintWithTimestamp(f"ADDRESS_URL = {addressUrl}")
PrintWithTimestamp(f"TIME_INTERVAL_SECONDS = {timeIntervalSeconds}")
PrintWithTimestamp("Number of unseal keys: " + str(len(unsealKeys)))

PrintDebug("UNSEAL_KEYS:")
for key in unsealKeys:
  PrintDebug(f"- {key}")
PrintWithTimestamp("If you do not see any output below, it means that the vault is contacted successfully and its unsealed")
PrintWithTimestamp(f"Vault will be contacted every {timeIntervalSeconds} seconds")
PrintWithTimestamp("Run with environment variable VU_DEBUG set to 1 for debug output")

while True:
  try:
    r = requests.get(f"{addressUrl}/v1/sys/seal-status").json()
    PrintDebug(f"status:{r}")
    if "sealed" in r:
      if r["sealed"] == True:
        PrintWithTimestamp("Detected sealed vault. Unsealing...")
        for key in unsealKeys:
          PrintDebug(f"key:{key}")
          r = requests.put(f"{addressUrl}/v1/sys/unseal", json = {"key":key}).json()
          PrintDebug(f"unseal:{r}")
        if r["sealed"] == True:
          PrintWithTimestamp("something went wrong, failed to unseal. Check the keys")
          PrintWithTimestamp(r)
          sys.exit(2)
        else:
          PrintWithTimestamp("Unsealed successfully")
    else:
      PrintWithTimestamp("Error: cannot find 'sealed' in returned json")
      pprint(r)
  except Exception as e:
    PrintWithTimestamp(f"Exception:{e}")
    PrintWithTimestamp(type(e))
    pprint(vars(e))
  time.sleep(timeIntervalSeconds)  
\$\endgroup\$
4
\$\begingroup\$

I will try to not repeat much of @Josay's answer since it was written while I was poking with your code. But the key points still applies:

  • The logging module is your friend
  • Follow PEP8 conventions
  • Don't compare truth values using == True
  • Don't repeat yourself

Speaking about the logging module, I simplified much of its use by removing the use of pprint. But you could add it back if you wish, here are some recipe. Keep in mind, though, that using logging functions like logger.debug is meant to minimise the runtime needed to build the message string, so using f-strings in this context (which is rather nice otherwise) would lead to unwanted overhead in case of "non-debug" runs. Oh, and, by the way, sys.stdout.write('…\n') is better understood as print('…').

You also happen to open, read and parse the configuration file thrice (once for each setting), this is a waste of resources. Read it once, use its values to update the os.environ dictionnary and search your settings in the resulting dictionnary.

The same kind of things happen when you try to unseal the vault. You use each key and then, check if the vault has successfully been unsealed. You can check after each key to avoid sending extraneous requests once the vault is unsealed. You can also use the for .. else construct to check that no keys unsealed the vault.

Proposed improvements:

import os
import sys
import json
import time
import logging

import requests


VERSION = '0.1'
ENV_PREFIX = 'VU_'
CONFIG_FILENAME = 'vault-unseal.json'


def configure_logger():
    logging.basicConfig(stream=sys.stdout, format='[%(asctime)s] %(message)s')
    logger = logging.getLogger('vault-unseal')
    logger.setLevel(logging.DEBUG if 'VU_DEBUG' in os.environ else logging.INFO)
    return logger


def read_configuration_file(filename=CONFIG_FILENAME):
    try:
        with open(filename) as f:
            return json.load(f)
    except (OSError, ValueError):
        return {}


def read_setting(setting_name, parameters, logger):
    if setting_name.endswith('*'):
        prefix = setting_name[:-1]
        setting = {key: value for key, value in parameters.items() if key.startswith(prefix)}
    else:
        setting = parameters.get(setting_name)

    if not setting:
        message = 'Error: %s is not specified. Specify it either in %s or as environment variable prefixed with %s'
        logger.error(message, setting_name, CONFIG_FILENAME, ENV_PREFIX)
        sys.exit(1)

    return setting


def unseal(base_url, time_interval, unseal_keys, logger):
    url = f'{base_url}/v1/sys/seal-status'
    unseal_url = f'{base_url}/v1/sys/unseal'

    while True:
        try:
            r = requests.get(url).json()
            logger.debug('status: %s', r)

            try:
                sealed = r['sealed']
            except KeyError:
                logger.error('Error: cannot find \'sealed\' in returned JSON\n%s', r)
            else:
                if sealed:
                    logger.info('Detected sealed vault. Unsealing…')

                    for key_name, key_value in unseal_keys.items():
                        logger.debug('Using key %s (%s)', key_name, key_value)
                        r = requests.put(unseal_url, json={'key': key_value}).json()
                        if r['sealed']:
                            logger.debug('Unseal result: %s', r)
                        else:
                            logger.info('Unsealed successfully')
                            break
                    else:
                        logger.error('Something went wrong, failed to unseal. Check the keys.\n%s', r)
                        sys.exit(2)
        except Exception:
            logger.exception('An exception occured:')
        time.sleep(time_interval)


def main():
    logger = configure_logger()
    logger.info('vault-unseal.py version %s', VERSION)
    logger.debug('Dumping environment block: \n%s', os.environ)

    # Filter environment variables of interest
    settings = {key[len(ENV_PREFIX):]: value for key, value in os.environ.items() if key.startswith(ENV_PREFIX)}
    # Update values with those found in configuration file
    settings.update(read_configuration_file())

    # Retrieve required parameters
    address_url = read_setting('ADDRESS_URL', settings, logger)
    time_interval = int(read_setting('TIME_INTERVAL_SECONDS', settings, logger))
    unseal_keys = read_setting('UNSEAL_KEY_*', settings, logger)

    logger.info('ADDRESS_URL = %s', address_url)
    logger.info('TIME_INTERVAL_SECONDS = %d', time_interval)
    logger.info('Number of unseal keys: %d', len(unseal_keys))
    logger.debug('UNSEAL_KEYS:')
    for key, value in unseal_keys.items():
        logger.debug('- %s: %s', key, value)

    logger.info('If you do not see any output below, it means that the vault is contacted successfully and its unsealed')
    logger.info('Vault will be contacted every %d seconds', time_interval)
    logger.info('Run with environment variable VU_DEBUG set for debug output')

    unseal(address_url, time_interval, unseal_keys, logger)


if __name__ == '__main__':
    main()
\$\endgroup\$
  • \$\begingroup\$ As I commented in the other answer I do not quite get the advice about not using == True. I want detect a case, where json part is looking like "sealed" :True. Note that "sealed" :"thingamabob" will also evaluate to True when compared without == True and that's something I like to avoid. Any comments? Thank you for your answer, I find it useful. \$\endgroup\$ – Andrew Savinykh Jan 4 '18 at 20:58
  • \$\begingroup\$ @AndrewSavinykh The "don't use == True" part comes mostly from Pep8 (scroll a little bit above) and is mostly a matter of readability. The same as you don't say "if the weather is fine is true, then whatever", you don't compare truthiness using if fine_weather == True: whatever(). Now as regard to the fact that non-empty strings also evaluate to true, I'd just avoid API that could return booleans or strings for the same result like the plague. \$\endgroup\$ – Mathias Ettinger Jan 4 '18 at 21:13
  • \$\begingroup\$ One more question. Without pprint the dictionary is dumped as one long string, which is difficult to read. Any suggestions? \$\endgroup\$ – Andrew Savinykh Jan 4 '18 at 21:20
  • \$\begingroup\$ @AndrewSavinykh check the use of pformat in the provided link. \$\endgroup\$ – Mathias Ettinger Jan 4 '18 at 21:25
  • 1
    \$\begingroup\$ Indeed. If only people read the entirety of the answer with all the links the first time around they would not need to ask stupid questions ;). Apologies, my mistake. Thank you, again. \$\endgroup\$ – Andrew Savinykh Jan 4 '18 at 21:32
4
\$\begingroup\$

Style

Your code does not follow the Python Style Guide (known as PEP 8). Along other things:

  • you should use 4-space indents

  • variable names should use snake_case

  • constant names should be UPPER_CASE

  • you should remove trailing whitespaces

I think the PrintXXX functions could have shorter names like log and debug. Also it could use the logging module.

Also, it is a good habit to move your code actually doing things behind an if __name__ == "__main__": guard.

Finally, it would be nice to have some documentation for the module and the different functions.

At this stage, the code looks like this (I haven't tested it).

import os
import sys
import json
import time
import datetime
from pprint import pprint
import requests

VERSION = "0.1"
ENV_PREFIX = "VU_"
CONFIG_FILE = "vault-unseal.json"


def debug_is_enabled():
    return ENV_PREFIX + "DEBUG" in os.environ


def log(string):
    timestamp = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S')
    sys.stdout.write(f"[{timestamp}] {string}\n")


def debug(string):
    if debug_is_enabled():
        log(f"DEBUG: {string}")


def print_param_error(name):
    log(f"Error: {name} is not specified. Specify it either in {CONFIG_FILE} file, or as enviroment variable prefixed with {ENV_PREFIX}")


def normalize_prefix(name, string):
    return name + string[len(ENV_PREFIX+name):]


def read_single_setting(name):
    result = None
    try:
        with open(CONFIG_FILE) as data_file:
            data = json.load(data_file)
        result = data[name]
    except (FileNotFoundError, KeyError):
        pass
    if ENV_PREFIX+name in os.environ:
        result = os.environ[ENV_PREFIX+name]
    return result


def read_multi_setting(name):
    result = None
    try:
        with open(CONFIG_FILE) as data_file:
            data = json.load(data_file)
        result = { key:data[key] for key in filter(lambda x: x.startswith(name), data.keys()) }
    except (FileNotFoundError):
        pass

    result = result if result else {}

    prefix = ENV_PREFIX + name
    env_result = { normalize_prefix(name,key):os.environ[key] for key in filter(lambda x: x.lower().startswith(prefix.lower()), list(os.environ.keys())) }

    final_result = list({**result,**env_result}.values())
    final_result = final_result if len(final_result) else None
    return final_result


def read_setting(name):
    if name.endswith("*"):
        return read_multi_setting(name[0:-1])
    else:
        return read_single_setting(name)



if __name__ == "__main__":
    log(f"vault-unseal.py VERSION {VERSION}")

    if debug_is_enabled():
        debug("Dumping environment block:")
        pprint(dict(os.environ))

    addressUrl = read_setting("ADDRESS_URL")
    timeIntervalSeconds = int(read_setting("TIME_INTERVAL_SECONDS"))
    unseal_keys = read_setting("UNSEAL_KEY_*")

    if not addressUrl:
        print_param_error("ADDRESS_URL")
        sys.exit(1)

    if not timeIntervalSeconds:
        print_param_error("TIME_INTERVAL_SECONDS")
        sys.exit(1)

    if not unseal_keys:
        print_param_error("UNSEAL_KEY_*")
        sys.exit(1)

    log(f"ADDRESS_URL = {addressUrl}")
    log(f"TIME_INTERVAL_SECONDS = {timeIntervalSeconds}")
    log("Number of unseal keys: " + str(len(unseal_keys)))

    debug("UNSEAL_KEYS:")
    for key in unseal_keys:
        debug(f"- {key}")
    log("If you do not see any output below, it means that the vault is contacted successfully and its unsealed")
    log(f"Vault will be contacted every {timeIntervalSeconds} seconds")
    log("Run with environment variable VU_DEBUG set to 1 for debug output")

    while True:
        try:
            r = requests.get(f"{addressUrl}/v1/sys/seal-status").json()
            debug(f"status:{r}")
            if "sealed" in r:
                if r["sealed"] == True:
                    log("Detected sealed vault. Unsealing...")
                    for key in unseal_keys:
                        debug(f"key:{key}")
                        r = requests.put(f"{addressUrl}/v1/sys/unseal", json = {"key":key}).json()
                        debug(f"unseal:{r}")
                    if r["sealed"] == True:
                        log("something went wrong, failed to unseal. Check the keys")
                        log(r)
                        sys.exit(2)
                    else:
                        log("Unsealed successfully")
            else:
                log("Error: cannot find 'sealed' in returned json")
                pprint(r)
        except Exception as e:
            log(f"Exception:{e}")
            log(type(e))
            pprint(vars(e))
        time.sleep(timeIntervalSeconds)

Do not repeat yourself

The check if "VU_DEBUG" in os.environ appears in multiple places. Maybe it would be easier to move this into a debug_is_enabled() function. Also, this could reuse the ENV_PREFIX constants like the other parts of the code:

def debug_is_enabled():
    return ENV_PREFIX + "DEBUG" in os.environ

Choose a wise default value

In read_multi_setting, you have:

result = None
// code_assigning_or_not_a_new_dict_to_result
result = result if result else {}

The last line only changes result if it is still None. It would be more straight-forward to write:

result = {}
// code_assigning_or_not_a_new_dict_to_result

Do not perform useless operations

In read_single_setting, you parse a file to get a value that will be overriden by environment variables. Maybe it'd be easier to perform the environ check first and parse the file only if nothing was found.

I'd write somethiing like:

def read_single_setting(name):
    if ENV_PREFIX+name in os.environ:
        return os.environ[ENV_PREFIX+name]
    try:
        with open(CONFIG_FILE) as data_file:
            data = json.load(data_file)
        return data[name]
    except (FileNotFoundError, KeyError):
        pass
    return None

Setting check

You retrieve 3 different settings and check if there are not false. I am not sure if this is intended but for timeIntervalSeconds, you perform a conversion to int before the check. I'd probably perform the conversion after the check (also, sleeping for 0 seconds seems to be valid).

Consistent return value

In read_multi_setting, final_result is initially a list. Then, there is:

final_result = final_result if len(final_result) else None

This seems confusing to me. Now, instead of final_result always being a list, either empty or not, it is now either a non-empty list or None. It make the function harder to use properly while making it more complicated. This could simply be:

return list({**result,**env_result}.values())

Simplify dict comprehension

You could call prefix.lower() before the loop.

You can use the if syntax in your comprehension in order not to call filter.

At this stage, you have:

prefix = (ENV_PREFIX + name).lower()
env_result = { normalize_prefix(name,key):os.environ[key] for key in list(os.environ.keys()) if key.lower().startswith(prefix) }

As far as I can tell, you don't need the conversion to list. Also, you can iterate directly over keys and values using items():

env_result = { normalize_prefix(name, key): value for key, value in os.environ.items() if key.lower().startswith(prefix) }

Comparison to True

You don't need to write, if value == True:, you can simply write if value:.

Use get to get value from dict with default value

Instead of checking if a value is in a dict and then get the value, you could do both in one go.

        if "sealed" in r:
            if r["sealed"]:

        if r.get("sealed", None):

At this stage, the code looks like:

import os
import sys
import json
import time
import datetime
from pprint import pprint
import requests

VERSION = "0.1"
ENV_PREFIX = "VU_"
CONFIG_FILE = "vault-unseal.json"


def debug_is_enabled():
    return ENV_PREFIX + "DEBUG" in os.environ


def log(string):
    timestamp = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S')
    sys.stdout.write(f"[{timestamp}] {string}\n")


def debug(string):
    if debug_is_enabled():
        log(f"DEBUG: {string}")


def print_param_error(name):
    log(f"Error: {name} is not specified. Specify it either in {CONFIG_FILE} file, or as enviroment variable prefixed with {ENV_PREFIX}")


def normalize_prefix(name, string):
    return name + string[len(ENV_PREFIX+name):]


def read_single_setting(name):
    if ENV_PREFIX+name in os.environ:
        return os.environ[ENV_PREFIX+name]
    try:
        with open(CONFIG_FILE) as data_file:
            data = json.load(data_file)
        return data[name]
    except (FileNotFoundError, KeyError):
        pass
    return None


def read_multi_setting(name):
    result = {}
    try:
        with open(CONFIG_FILE) as data_file:
            data = json.load(data_file)
        result = { key:data[key] for key in filter(lambda x: x.startswith(name), data.keys()) }
    except (FileNotFoundError):
        pass

    prefix = (ENV_PREFIX + name).lower()
    env_result = { normalize_prefix(name, key): value for key, value in os.environ.items() if key.lower().startswith(prefix) }

    return list({**result,**env_result}.values())


def read_setting(name):
    if name.endswith("*"):
        return read_multi_setting(name[0:-1])
    else:
        return read_single_setting(name)



if __name__ == "__main__":
    log(f"vault-unseal.py VERSION {VERSION}")

    if debug_is_enabled():
        debug("Dumping environment block:")
        pprint(dict(os.environ))

    addressUrl = read_setting("ADDRESS_URL")
    timeIntervalSeconds = int(read_setting("TIME_INTERVAL_SECONDS"))
    unseal_keys = read_setting("UNSEAL_KEY_*")

    if not addressUrl:
        print_param_error("ADDRESS_URL")
        sys.exit(1)

    if not timeIntervalSeconds:
        print_param_error("TIME_INTERVAL_SECONDS")
        sys.exit(1)

    if not unseal_keys:
        print_param_error("UNSEAL_KEY_*")
        sys.exit(1)

    log(f"ADDRESS_URL = {addressUrl}")
    log(f"TIME_INTERVAL_SECONDS = {timeIntervalSeconds}")
    log("Number of unseal keys: " + str(len(unseal_keys)))

    debug("UNSEAL_KEYS:")
    for key in unseal_keys:
        debug(f"- {key}")
    log("If you do not see any output below, it means that the vault is contacted successfully and its unsealed")
    log(f"Vault will be contacted every {timeIntervalSeconds} seconds")
    log("Run with environment variable VU_DEBUG set to 1 for debug output")

    while True:
        try:
            r = requests.get(f"{addressUrl}/v1/sys/seal-status").json()
            debug(f"status:{r}")
            if r.get("sealed", None):
                log("Detected sealed vault. Unsealing...")
                for key in unseal_keys:
                    debug(f"key:{key}")
                    r = requests.put(f"{addressUrl}/v1/sys/unseal", json = {"key":key}).json()
                    debug(f"unseal:{r}")
                if r["sealed"]:
                    log("something went wrong, failed to unseal. Check the keys")
                    log(r)
                    sys.exit(2)
                else:
                    log("Unsealed successfully")
            else:
                log("Error: cannot find 'sealed' in returned json")
                pprint(r)
        except Exception as e:
            log(f"Exception:{e}")
            log(type(e))
            pprint(vars(e))
        time.sleep(timeIntervalSeconds)
\$\endgroup\$
  • \$\begingroup\$ Thank you. r["sealed"] is a boolean so your suggestion regarding it does not work quite as intended. In particular r.get("sealed", None) condition will evaluate to False if the json value is False. This is not intended. I'm open to suggestions how to improve that part. \$\endgroup\$ – Andrew Savinykh Jan 4 '18 at 21:00

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.