CloudBender/cloudbender/jinja.py

278 lines
8.0 KiB
Python
Raw Permalink Normal View History

import os
2018-11-22 18:31:59 +00:00
import io
import gzip
import re
import base64
import yaml
import copy
import subprocess
import sys
import zlib
import jinja2
import python_minifier
2022-04-19 12:02:31 +00:00
import markupsafe
2019-04-18 16:30:50 +00:00
from jinja2.filters import make_attrgetter
from jinja2.runtime import Undefined
2018-11-22 18:31:59 +00:00
import logging
2019-04-18 16:30:50 +00:00
2018-11-22 18:31:59 +00:00
logger = logging.getLogger(__name__)
2019-02-05 17:48:29 +00:00
@jinja2.pass_context
2022-02-22 10:04:29 +00:00
def option(context, attribute, default_value="", source="options"):
"""Get attribute from options data structure, default_value otherwise"""
2019-04-18 16:30:50 +00:00
environment = context.environment
2022-02-22 10:04:29 +00:00
options = environment.globals["_config"][source]
2019-02-05 17:48:29 +00:00
2019-04-18 16:30:50 +00:00
if not attribute:
return default_value
2019-02-05 17:48:29 +00:00
2019-04-18 16:30:50 +00:00
try:
getter = make_attrgetter(environment, attribute)
value = getter(options)
2019-02-05 17:48:29 +00:00
2019-04-18 16:30:50 +00:00
if isinstance(value, Undefined):
return default_value
2019-04-18 16:30:50 +00:00
return value
2019-02-05 17:48:29 +00:00
2019-04-18 16:30:50 +00:00
except (jinja2.exceptions.UndefinedError):
return default_value
2018-11-22 18:31:59 +00:00
@jinja2.pass_context
def include_raw_gz(context, files=None, gz=True, remove_comments=False):
2018-11-22 18:31:59 +00:00
jenv = context.environment
2022-02-22 10:04:29 +00:00
output = ""
2020-07-16 22:56:09 +00:00
# For shell script we can even remove whitespaces so treat them individually
# sed -e '2,$ {/^ *$/d ; /^ *#/d ; /^[ \t] *#/d ; /*^/d ; s/^[ \t]*// ; s/*[ \t]$// ; s/ $//}'
2018-11-22 18:31:59 +00:00
for name in files:
2022-04-19 12:02:31 +00:00
output = output + markupsafe.Markup(jenv.loader.get_source(jenv, name)[0])
2018-11-22 18:31:59 +00:00
if remove_comments:
# Remove full line comments but not shebang
2022-02-22 10:04:29 +00:00
_re_comment = re.compile(r"^\s*#[^!]")
_re_blank = re.compile(r"^\s*$")
_re_keep = re.compile(r"^## template: jinja$")
stripped_output = ""
for curline in output.splitlines():
if re.match(_re_blank, curline):
continue
elif re.match(_re_keep, curline):
2022-02-22 10:04:29 +00:00
stripped_output = stripped_output + curline + "\n"
elif re.match(_re_comment, curline):
logger.debug("Removed {}".format(curline))
else:
2022-02-22 10:04:29 +00:00
stripped_output = stripped_output + curline + "\n"
output = stripped_output
2018-11-22 18:31:59 +00:00
if not gz:
2022-02-22 10:04:29 +00:00
return output
2018-11-22 18:31:59 +00:00
buf = io.BytesIO()
2022-02-22 10:04:29 +00:00
f = gzip.GzipFile(mode="w", fileobj=buf, mtime=0)
2018-11-22 18:31:59 +00:00
f.write(output.encode())
f.close()
# MaxSize is 21847
2022-02-22 10:04:29 +00:00
logger.info(
"Compressed user-data from {} to {}".format(len(output), len(buf.getvalue()))
)
return base64.b64encode(buf.getvalue()).decode("utf-8")
2018-11-22 18:31:59 +00:00
@jinja2.pass_context
2018-11-22 18:31:59 +00:00
def raise_helper(context, msg):
raise Exception(msg)
# Custom tests
2022-02-22 10:04:29 +00:00
def regex(value="", pattern="", ignorecase=False, match_type="search"):
"""Expose `re` as a boolean filter using the `search` method by default.
This is likely only useful for `search` and `match` which already
have their own filters.
"""
2018-11-22 18:31:59 +00:00
if ignorecase:
flags = re.I
else:
flags = 0
_re = re.compile(pattern, flags=flags)
2022-02-22 10:04:29 +00:00
if getattr(_re, match_type, "search")(value) is not None:
return True
2018-11-22 18:31:59 +00:00
return False
2022-02-22 10:04:29 +00:00
def match(value, pattern="", ignorecase=False):
"""Perform a `re.match` returning a boolean"""
return regex(value, pattern, ignorecase, "match")
2018-11-22 18:31:59 +00:00
2022-02-22 10:04:29 +00:00
def search(value, pattern="", ignorecase=False):
"""Perform a `re.search` returning a boolean"""
return regex(value, pattern, ignorecase, "search")
2018-11-22 18:31:59 +00:00
# Custom filters
2022-02-22 10:04:29 +00:00
def sub(value="", pattern="", replace="", ignorecase=False):
2018-11-22 18:31:59 +00:00
if ignorecase:
flags = re.I
else:
flags = 0
return re.sub(pattern, replace, value, flags=flags)
2018-11-22 18:31:59 +00:00
def pyminify(source):
minified = python_minifier.awslambda(source, filename=None, entrypoint=None)
gz_source = gz_pack(minified)
logger.info(
"Compressed python code from {} to {}".format(len(source), len(gz_source))
2022-02-22 10:04:29 +00:00
)
return gz_source
2018-11-22 18:31:59 +00:00
# From pyminifier
def gz_pack(source):
"""
Returns 'source' as a gzip-compressed, self-extracting python script.
2018-11-22 18:31:59 +00:00
.. note::
2018-11-22 18:31:59 +00:00
This method uses up more space than the zip_pack method but it has the
advantage in that the resulting .py file can still be imported into a
python program.
"""
out = ""
# Preserve shebangs (don't care about encodings for this)
first_line = source.split("\n")[0]
if re.compile("^#!.*$").match(first_line):
if first_line.rstrip().endswith("python"):
first_line = first_line.rstrip()
first_line += "3"
out = first_line + "\n"
compressed_source = zlib.compress(source.encode("utf-8"))
out += "import zlib, base64\n"
out += "exec(zlib.decompress(base64.b64decode('"
out += base64.b64encode(compressed_source).decode("utf-8")
out += "')))\n"
return out
2018-11-22 18:31:59 +00:00
2019-04-18 16:30:50 +00:00
def inline_yaml(block):
return yaml.safe_load(block)
2018-11-22 18:31:59 +00:00
def JinjaEnv(template_locations=[]):
2022-01-24 11:01:50 +00:00
LoggingUndefined = jinja2.make_logging_undefined(logger=logger, base=Undefined)
2022-02-22 10:04:29 +00:00
jenv = jinja2.Environment(
trim_blocks=True,
lstrip_blocks=True,
undefined=LoggingUndefined,
extensions=["jinja2.ext.loopcontrols", "jinja2.ext.do"],
)
2018-11-22 18:31:59 +00:00
if template_locations:
jinja_loaders = []
for _dir in template_locations:
jinja_loaders.append(jinja2.FileSystemLoader(str(_dir)))
jenv.loader = jinja2.ChoiceLoader(jinja_loaders)
else:
jenv.loader = jinja2.BaseLoader()
2018-11-22 18:31:59 +00:00
2022-02-22 10:04:29 +00:00
jenv.globals["include_raw"] = include_raw_gz
jenv.globals["raise"] = raise_helper
jenv.globals["option"] = option
2018-11-22 18:31:59 +00:00
2022-02-22 10:04:29 +00:00
jenv.filters["sub"] = sub
jenv.filters["pyminify"] = pyminify
jenv.filters["inline_yaml"] = inline_yaml
2018-11-22 18:31:59 +00:00
2022-02-22 10:04:29 +00:00
jenv.tests["match"] = match
jenv.tests["regex"] = regex
jenv.tests["search"] = search
2018-11-22 18:31:59 +00:00
return jenv
def render_docs(docs, outputs):
jenv = jinja2.Environment(undefined=jinja2.ChainableUndefined)
return jenv.from_string(docs).render(outputs)
def read_config_file(path, variables={}):
2022-02-22 10:04:29 +00:00
"""reads yaml config file, passes it through jinja and returns data structre
2022-02-22 10:04:29 +00:00
- OS ENV are available as {{ ENV.<VAR> }}
- variables defined in parent configs are available as {{ <VAR> }}
"""
jinja_variables = copy.deepcopy(variables)
2022-02-22 10:04:29 +00:00
jinja_variables["ENV"] = os.environ
if path.exists():
logger.debug("Reading config file: {}".format(path))
# First check for sops being present
try:
jenv = jinja2.Environment(
enable_async=True,
auto_reload=False,
loader=jinja2.FunctionLoader(_sops_loader),
undefined=jinja2.StrictUndefined,
2022-02-22 10:04:29 +00:00
extensions=["jinja2.ext.loopcontrols"],
)
template = jenv.get_template(str(path))
rendered_template = template.render(jinja_variables)
data = yaml.safe_load(rendered_template)
if data:
return data
except Exception as e:
logger.exception("Error reading config file: {} ({})".format(path, e))
sys.exit(1)
return {}
def _sops_loader(path):
2022-02-22 10:04:29 +00:00
"""Tries to loads yaml file
If "sops" key is detected the file is piped through sops before returned
"""
2022-02-22 10:04:29 +00:00
with open(path, "r") as f:
config_raw = f.read()
data = yaml.safe_load(config_raw)
2022-02-22 10:04:29 +00:00
if data and "sops" in data and "DISABLE_SOPS" not in os.environ:
try:
2022-02-22 10:04:29 +00:00
result = subprocess.run(
[
"sops",
"--input-type",
"yaml",
"--output-type",
"yaml",
"--decrypt",
"/dev/stdin",
],
stdout=subprocess.PIPE,
input=config_raw.encode("utf-8"),
env=dict(os.environ, **{"AWS_SDK_LOAD_CONFIG": "1"}),
)
except FileNotFoundError:
2022-02-22 10:04:29 +00:00
logger.exception(
"SOPS encrypted config {}, but unable to find sops binary! Try eg: https://github.com/mozilla/sops/releases/download/v3.5.0/sops-v3.5.0.linux".format(
path
)
)
sys.exit(1)
2022-02-22 10:04:29 +00:00
return result.stdout.decode("utf-8")
else:
return config_raw