87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
import re
|
|
import yaml
|
|
|
|
|
|
__all__ = ['safe_dump', 'SafeLoader']
|
|
|
|
|
|
class SafeStringDumper(yaml.SafeDumper):
|
|
|
|
def represent_data(self, value):
|
|
if isinstance(value, str):
|
|
return self.represent_scalar('!unsafe', value)
|
|
return super(SafeStringDumper, self).represent_data(value)
|
|
|
|
|
|
class SafeLoader(yaml.Loader):
|
|
|
|
def construct_yaml_unsafe(self, node):
|
|
class UnsafeText(str):
|
|
__UNSAFE__ = True
|
|
node = UnsafeText(self.construct_scalar(node))
|
|
return node
|
|
|
|
|
|
SafeLoader.add_constructor(
|
|
u'!unsafe',
|
|
SafeLoader.construct_yaml_unsafe
|
|
)
|
|
|
|
|
|
def safe_dump(x, safe_dict=None):
|
|
"""
|
|
Used to serialize an extra_vars dict to YAML
|
|
|
|
By default, extra vars are marked as `!unsafe` in the generated yaml
|
|
_unless_ they've been deemed "trusted" (meaning, they likely were set/added
|
|
by a user with a high level of privilege).
|
|
|
|
This function allows you to pass in a trusted `safe_dict` to allow
|
|
certain extra vars so that they are _not_ marked as `!unsafe` in the
|
|
resulting YAML. Anything _not_ in this dict will automatically be
|
|
`!unsafe`.
|
|
|
|
safe_dump({'a': 'b', 'c': 'd'}) ->
|
|
!unsafe 'a': !unsafe 'b'
|
|
!unsafe 'c': !unsafe 'd'
|
|
|
|
safe_dump({'a': 'b', 'c': 'd'}, safe_dict={'a': 'b'})
|
|
a: b
|
|
!unsafe 'c': !unsafe 'd'
|
|
"""
|
|
if isinstance(x, dict):
|
|
yamls = []
|
|
safe_dict = safe_dict or {}
|
|
|
|
# Compare the top level keys so that we can find values that have
|
|
# equality matches (and consider those branches safe)
|
|
for k, v in x.items():
|
|
dumper = yaml.SafeDumper
|
|
if k not in safe_dict or safe_dict.get(k) != v:
|
|
dumper = SafeStringDumper
|
|
yamls.append(yaml.dump_all(
|
|
[{k: v}],
|
|
None,
|
|
Dumper=dumper,
|
|
default_flow_style=False,
|
|
))
|
|
return ''.join(yamls)
|
|
else:
|
|
return yaml.dump_all([x], None, Dumper=SafeStringDumper, default_flow_style=False)
|
|
|
|
|
|
def sanitize_jinja(arg):
|
|
"""
|
|
For some string, prevent usage of Jinja-like flags
|
|
"""
|
|
if isinstance(arg, str):
|
|
# If the argument looks like it contains Jinja expressions
|
|
# {{ x }} ...
|
|
if re.search(r'\{\{[^}]+}}', arg) is not None:
|
|
raise ValueError('Inline Jinja variables are not allowed.')
|
|
# If the argument looks like it contains Jinja statements/control flow...
|
|
# {% if x.foo() %} ...
|
|
if re.search(r'\{%[^%]+%}', arg) is not None:
|
|
raise ValueError('Inline Jinja variables are not allowed.')
|
|
return arg
|