diff --git a/pywb/utils/loaders.py b/pywb/utils/loaders.py index 6d721de3..f2bac05d 100644 --- a/pywb/utils/loaders.py +++ b/pywb/utils/loaders.py @@ -11,10 +11,14 @@ import requests import yaml import six -from six.moves.urllib.parse import unquote_plus, urlsplit +from six.moves.urllib.parse import unquote_plus, urlsplit, urlencode import time import pkgutil +import base64 +import yaml +import cgi +import re from io import open, BytesIO from warcio.limitreader import LimitReader @@ -30,7 +34,21 @@ except ImportError: # pragma: no cover s3_avail = False -# ================================================================= +# ============================================================================ +def init_yaml_env_vars(): + env_rx = re.compile(r'\$\{[^}]+\}') + + yaml.add_implicit_resolver('!envvar', env_rx) + + def envvar_constructor(loader, node): + value = loader.construct_scalar(node) + value = os.path.expandvars(value) + return value + + yaml.add_constructor('!envvar', envvar_constructor) + + +# ============================================================================ def load_py_name(string): import importlib @@ -468,3 +486,5 @@ class HMACCookieMaker(object): # ============================================================================ BlockLoader.init_default_loaders() + +init_yaml_env_vars() diff --git a/pywb/utils/test/test_loaders.py b/pywb/utils/test/test_loaders.py index 819390b3..9372f822 100644 --- a/pywb/utils/test/test_loaders.py +++ b/pywb/utils/test/test_loaders.py @@ -80,6 +80,7 @@ import six from six import StringIO from io import BytesIO import requests +import yaml from pywb.utils.loaders import BlockLoader, HMACCookieMaker, to_file_url from pywb.utils.loaders import extract_client_cookie @@ -168,6 +169,28 @@ def test_err_unknown_loader(): #IOError: No Loader for type: foo + +def test_yaml_resolve_env(): + os.environ['PYWB_PATH'] = './test' + os.environ['PYWB_FOO'] = 'bar' + + config = """\ +collection: + coll: + index: ${PYWB_PATH}/index + archive: ${PYWB_PATH}/archive/${PYWB_FOO} + other: ${PYWB_NOT}/archive/${PYWB_FOO} +""" + + config_data = yaml.load(config) + + assert config_data['collection']['coll']['index'] == './test/index' + assert config_data['collection']['coll']['archive'] == './test/archive/bar' + assert config_data['collection']['coll']['other'] == '${PYWB_NOT}/archive/bar' + + del os.environ['PYWB_PATH'] + del os.environ['PYWB_FOO'] + def print_str(string): return string.decode('utf-8') if six.PY3 else string