def testListScopeVarsByteString(self):
interpolator = interpolation.Interpolator(
b"%%foo.A%% %%foo.B%% %%foo.C%%")
self.assertEqual(interpolator.ScopeVars(sid("foo")), {
vid("A"),
vid("B"),
vid("C"),
})
def testInterpolationHappensSimultaneously(self):
interpolator = interpolation.Interpolator("%%foo%% %%bar.baz%% %%quux%%")
interpolator.BindVar(vid("foo"), "%%bar.baz%%")
interpolator.BindVar(vid("quux"), "%%foo%%")
interpolator.BindScope(sid("bar"), {"baz": "%%foo%% %%quux%%"})
strings = list(interpolator.Interpolate())
self.assertEqual(strings, ["%%bar.baz%% %%foo%% %%quux%% %%foo%%"])
def testBindScopeSingleUnicodeString(self):
interpolator = interpolation.Interpolator("%%foo.bar%% %%foo.baz%%")
interpolator.BindScope(sid("foo"), {
vid("bar"): "quux",
vid("baz"): "norf"
})
strings = list(interpolator.Interpolate())
self.assertEqual(strings, ["quux norf"])
def testBindVarTwoIntegersToByteString(self):
interpolator = interpolation.Interpolator(b"%%foo%%%%bar%%")
interpolator.BindVar(vid("foo"), 1)
interpolator.BindVar(vid("foo"), 2)
interpolator.BindVar(vid("bar"), 3)
interpolator.BindVar(vid("bar"), 4)
strings = list(interpolator.Interpolate())
self.assertCountEqual(strings, [b"13", b"14", b"23", b"24"])
def testDuplicatedVars(self):
interpolator = interpolation.Interpolator("%%foo%%%%bar%%%%foo%%%%bar%%")
interpolator.BindVar(vid("foo"), 1)
interpolator.BindVar(vid("foo"), 2)
interpolator.BindVar(vid("bar"), 3)
interpolator.BindVar(vid("bar"), 4)
strings = list(interpolator.Interpolate())
self.assertCountEqual(strings, ["1313", "1414", "2323", "2424"])
def testBindScopeSingleByteString(self):
interpolator = interpolation.Interpolator(
b"%%foo.A%% %%foo.B%% %%foo.C%%")
interpolator.BindScope(sid("foo"), {
vid("A"): 1,
vid("B"): 2,
vid("C"): 3
})
strings = list(interpolator.Interpolate())
self.assertEqual(strings, [b"1 2 3"])
def testBindScopeMultipleUnicodeString(self):
pattern = "%%foo.bar%%$%%quux.norf%%$%%foo.baz%%$%%quux.thud%%"
interpolator = interpolation.Interpolator(pattern)
interpolator.BindScope(sid("foo"), {vid("bar"): 1, vid("baz"): 2})
interpolator.BindScope(sid("foo"), {vid("bar"): 3, vid("baz"): 4})
interpolator.BindScope(sid("quux"), {vid("norf"): 5, vid("thud"): 6})
interpolator.BindScope(sid("quux"), {vid("norf"): 7, vid("thud"): 8})
strings = list(interpolator.Interpolate())
self.assertCountEqual(strings, ["1$5$2$6", "1$7$2$8", "3$5$4$6", "3$7$4$8"])
def testBindScopeAndVarUnicodeString(self):
pattern = "%%foo.bar%%|%%quux%%|%%foo.baz%%|%%norf.thud%%|%%norf.blargh%%"
interpolator = interpolation.Interpolator(pattern)
interpolator.BindVar(vid("quux"), 1)
interpolator.BindVar(vid("quux"), 2)
interpolator.BindScope(sid("foo"), {vid("bar"): 3, vid("baz"): 4})
interpolator.BindScope(sid("foo"), {vid("bar"): 5, vid("baz"): 6})
interpolator.BindScope(sid("norf"), {vid("thud"): 7, vid("blargh"): 8})
interpolator.BindScope(sid("norf"), {vid("thud"): 9, vid("blargh"): 0})
strings = list(interpolator.Interpolate())
self.assertCountEqual(strings, [
"3|1|4|7|8",
"3|1|4|9|0",
"3|2|4|7|8",
"3|2|4|9|0",
"5|1|6|7|8",
"5|1|6|9|0",
"5|2|6|7|8",
"5|2|6|9|0",
])
def testBindVarKeyError(self):
interpolator = interpolation.Interpolator("%%foo%%")
with self.assertRaises(KeyError):
interpolator.BindVar(vid("bar"), 42)
def testBindVarWeirdByteString(self):
interpolator = interpolation.Interpolator(b"\xff %%foo%% \xff")
interpolator.BindVar(vid("foo"), 42)
strings = list(interpolator.Interpolate())
self.assertEqual(strings, [b"\xff 42 \xff"])
def testBindVarWeirdUnicodeString(self):
interpolator = interpolation.Interpolator("❄ %%foo%% 🌊")
interpolator.BindVar(vid("foo"), "🎄")
strings = list(interpolator.Interpolate())
self.assertEqual(strings, ["❄ 🎄 🌊"])
def testBindVarSimpleByteString(self):
interpolator = interpolation.Interpolator(b"foo %%bar%% baz")
interpolator.BindVar(vid("bar"), "quux")
strings = list(interpolator.Interpolate())
self.assertEqual(strings, [b"foo quux baz"])
def testListScopeVarsUnicodeString(self):
interpolator = interpolation.Interpolator(
"%%foo.A%% %%foo.B%% %%bar.C%%")
self.assertEqual(interpolator.ScopeVars(sid("foo")),
{vid("A"), vid("B")})
self.assertEqual(interpolator.ScopeVars(sid("bar")), {vid("C")})
def testListScopeByteString(self):
interpolator = interpolation.Interpolator(b"%%foo.bar%% %%quux.norf%%")
self.assertEqual(interpolator.Scopes(), {sid("foo"), sid("quux")})
def testListVarsByteString(self):
interpolator = interpolation.Interpolator(
b"%%foo%%%%bar%%%%bar%%%%baz%%")
self.assertEqual(
interpolator.Vars(),
{vid("foo"), vid("bar"), vid("baz")})
def InterpolateKbAttributes(pattern, knowledge_base, ignore_errors=False):
"""Interpolate all knowledgebase attributes in pattern.
Args:
pattern: A string with potential interpolation markers. For example:
"/home/%%users.username%%/Downloads/"
knowledge_base: The knowledge_base to interpolate parameters from.
ignore_errors: Set this to true to log errors instead of raising.
Raises:
KnowledgeBaseInterpolationError: If there are any interpolation problems
and the `ignore_errors` flag is not set.
Yields:
All unique strings generated by expanding the pattern.
"""
# TODO(hanuszczak): Control flow feels a bit awkward here because of error
# handling that tries not to break any functionality. With the new utilities
# it should be possible to improve the code, changing the behaviour to a more
# sane one.
missing_attr_names = set()
interpolator = interpolation.Interpolator(pattern)
for vid in interpolator.Vars():
attr_name = str(vid).lower()
value = getattr(knowledge_base, attr_name, None)
if not value:
missing_attr_names.add(attr_name)
continue
interpolator.BindVar(attr_name, value)
for scope_id in interpolator.Scopes():
scope_name = str(scope_id).lower()
kb_structs = getattr(knowledge_base, scope_name, None)
if not kb_structs:
missing_attr_names.add(scope_name)
continue
bound = False
for kb_struct in kb_structs:
bindings = {}
for var_id in interpolator.ScopeVars(scope_id):
attr_name = str(var_id).lower()
value = getattr(kb_struct, attr_name, None)
if value:
bindings[var_id] = value
if len(bindings) == len(interpolator.ScopeVars(scope_id)):
interpolator.BindScope(scope_id, bindings)
bound = True
if not bound:
missing_attr_names.add(scope_name)
if missing_attr_names:
message = "Some attributes could not be located in the knowledgebase: {}"
message = message.format(", ".join(missing_attr_names))
if not ignore_errors:
raise KnowledgeBaseInterpolationError(message)
else:
logging.error(message)
return
for result in interpolator.Interpolate():
yield result
