def test_tc_argcheck(self):
"""test_tc_argcheck"""
tcex = tcex_dummy()
assert tc_argcheck(self, 'b', required=False) is None
assert tc_argcheck(self, 'foo', required=True, tcex=tcex) == 'bla'
tc_argcheck(self, 'b', required=True, tcex=tcex)
assert tcex.exitrc == 1
assert tcex.exitmsg == 'Invalid value for B. Value is required.'
def write_one(self, name, value, output_type):
"""Write one output"""
# JSON-ify any array outputs that are compound for the text outputs
if output_type in ('StringArray', 'BinaryArray'):
result = []
for entry in value:
if isinstance(entry, (list, dict, tuple)):
entry = json.dumps(entry, sort_keys=True)
result.append(entry)
value = result
if output_type in ('String', 'Binary') and isinstance(
value, (dict, list, tuple)):
value = json.dumps(value, sort_keys=True)
self.tcex.log.debug(f'Write: {name}!{output_type}={value!r}')
try:
self.tcex.playbook.create_output(name, value, output_type)
except Exception as e:
return_none_on_failure = tc_argcheck(self.tcex.args,
'return_none_on_failure',
tcex=self.tcex)
self.tcex.log.error(f'Error creating output {name}: {e}')
self.errors.append(f'Error creating output {name}: {e}')
if not return_none_on_failure:
self.tcex.playbook.exit(
1, f'Unable to create output {name}: {e}')
def handle_exception(self, exception, force_exit=False):
"""handle exception"""
return_none_on_failure = tc_argcheck(self.tcex.args,
'return_none_on_failure',
tcex=self.tcex)
self.errors.append(str(exception))
self.tcex.log.error(str(exception))
tb = sys.exc_info()[2]
while tb and tb.tb_next:
tb = tb.tb_next
if tb:
name = tb.tb_frame.f_code.co_name
args = []
for arg in tb.tb_frame.f_code.co_varnames:
value = repr(tb.tb_frame.f_locals.get(arg))
if len(value) > 12:
value = value[:12] + '...'
if value.startswith('"'):
value += '"'
elif value.startswith("'"):
value += "'"
elif value.startswith('{'):
value += '}'
args.append(f'{arg}={value}')
args = ', '.join(args)
self.tcex.log.debug(f'Error during {name}({args})')
error = traceback.format_exc()
self.tcex.log.debug(error)
if force_exit or not return_none_on_failure:
self.tcex.playbook.exit(1, str(exception))
def setup_vars(self, varname):
"""Setup initial variables"""
var_list = tc_argcheck(self.tcex.args, varname, tcex=self.tcex)
if var_list:
var_list = self.tcex.playbook.read(var_list,
array=True,
embedded=False)
if not var_list:
var_list = []
for var in var_list:
name = var['key']
if not identifierRE.match(name):
self.handle_exception(
f'Invalid variable {name}. Variables must start with a letter and '
'contain only letters, numbers, and underscore.',
force_exit=True,
)
value = var['value']
self.tcex.log.debug(f'Setting {name} from {value}')
try:
if value:
value = self.engine.eval(value)
except Exception as e:
value = self.handle_exception(
f'Error evaluating {name}: {e}. For string literals, enclose them in quotes.'
)
self.tcex.log.debug(f'... {name} = {value!r}')
self.engine.set(name, value)
def setup_loops(self, varname):
""" Read the Key Value Array in, and setup iter_control """
loopvars = tc_argcheck(self.tcex.args,
varname,
required=True,
tcex=self.tcex)
loopvars = self.tcex.playbook.read(loopvars,
array=True,
embedded=False)
self.tcex.log.debug(f'Loopvars are {loopvars}')
sizes = []
for var in loopvars:
key = var['key']
if not identifierRE.match(key):
self.handle_exception(
f'Invalid loop variable {key}. Variables must start with a letter '
'and contain only letters, numbers, and underscore.',
force_exit=True,
)
value = var['value']
self.tcex.log.debug(f'Retrieving values for {key} from {value!r}')
if isinstance(value, str):
try:
value = self.engine.eval(value)
except Exception as e:
value = self.handle_exception(
f'Unable to get values for loop variable {key}: {e}')
if not isinstance(value, (list, tuple)):
value = (value, )
self.tcex.log.debug(f'Loop variable {key} is {value!r}')
sizes.append((len(value), key, value))
sizes.sort() # trackers from smallest to largest
iter_control = []
lastsize = None
tracker = None
for s, name, value in sizes:
if s != lastsize:
if tracker:
iter_control.append(tracker)
lastsize = s
tracker = IterTracker()
tracker.add(name, value)
if tracker:
iter_control.append(tracker)
return iter_control
def setup(self):
"""Perform prep/startup operations."""
self.nameservers = tc_argcheck(
self.tcex.rargs, 'dns_servers', label='DNS Servers', required=True, tcex=self.tcex
)
rate_limit = tc_argcheck(
self.tcex.rargs, 'rate_limit', required=True, types=int, default=150, tcex=self.tcex
)
self.throttle = Throttle(rate_limit)
self.transform_ptr = tc_argcheck(
self.tcex.rargs, 'transform_ptr', default=True, types=bool, tcex=self.tcex
)
if isinstance(self.nameservers, str):
self.nameservers = self.nameservers.split(',')
if not isinstance(self.nameservers, list):
self.nameservers = [self.nameservers]
self.nameservers = [x.strip() for x in self.nameservers]
self.add_output('dns.action', self.tcex.rargs.tc_action)
def evaluate_many_with_loop(self):
"""Evaluate many expressions with loop variables"""
self.setup_vars('variables')
iter_control = self.setup_loops('loop_variables')
exprs = tc_argcheck(self.tcex.args,
'loop_expressions',
required=True,
tcex=self.tcex)
if exprs:
exprs = self.tcex.playbook.read(exprs, embedded=False)
loop_exprs = {}
for expr in exprs:
key = expr['key']
value = expr['value']
loop_exprs[key] = value
self.outloop = self.loop_over_iter(iter_control, loop_exprs)
for key, value in self.outloop.items():
self.engine.set(key, value)
self.setup_outputs('additional_outputs', required=False)
self.setup_outputs('binary_outputs',
required=False,
output_type='Binary')
self.setup_outputs('binary_array_outputs',
required=False,
output_type='BinaryArray')
self.setup_outputs('kv_outputs',
required=False,
output_type='KeyValue')
self.setup_outputs('kv_array_outputs',
required=False,
output_type='KeyValueArray')
self.setup_outputs('tce_outputs',
required=False,
output_type='TCEntity')
self.setup_outputs('tce_array_outputs',
required=False,
output_type='TCEntityArray')
self.setup_outputs('tcee_outputs',
required=False,
output_type='TCEnhancedEntity')
self.setup_outputs('tcee_array_outputs',
required=False,
output_type='TCEnhancedEntityArray')
def evaluate_in_loop(self):
"""Evaluate an expression in a loop"""
iter_control = self.setup_loops('loop_variables')
expr = tc_argcheck(self.tcex.args,
'loop_expression',
required=True,
tcex=self.tcex,
label='Expression')
expr = self.tcex.playbook.read(expr, embedded=False)
self.expression = self.tcex.playbook.read(expr, embedded=True)
self.tcex.log.debug(f'Expression is {expr}')
outdict = self.loop_over_iter(iter_control, expr)
self.output = outdict.get('output')
def Evaluate(self):
"""Evaluate an expression"""
expr = tc_argcheck(self.tcex.args,
'expression',
required=True,
tcex=self.tcex)
self.expression = self.tcex.playbook.read(expr, embedded=True)
expr = self.tcex.playbook.read(expr, embedded=False)
self.tcex.log.info(f'Evaluating expression {expr!r}')
try:
result = self.engine.eval(expr)
except Exception as e:
result = self.handle_exception(
f'Evaluation of "{expr}" failed: {e}')
if isinstance(result, list):
self.output.extend(result)
else:
self.output.append(result)
def setup_outputs(self, varname, required=True, output_type='String'):
""" Setup output varables from expressions """
out_list = tc_argcheck(self.tcex.args,
varname,
required=required,
tcex=self.tcex)
out_list = self.tcex.playbook.read(out_list,
array=True,
embedded=False)
if not out_list:
return
for out in out_list:
name = out['key']
value = out['value']
self.tcex.log.debug(f'Setting {name} from {value}')
try:
if value:
value = self.engine.eval(value)
except Exception as e:
value = self.handle_exception(
f'Error evaluating {name}: {e}. For string literals, enclose them in quotes.'
)
self.engine.set(name, value) # allow outputs to see prior outputs
if isinstance(
value,
(list, tuple, dict)) and output_type in ('String', 'Binary'):
value = json.dumps(value, sort_keys=True, ensure_ascii=False)
if output_type.endswith('Array') and not isinstance(
value, (list, tuple)):
self.tcex.log.debug('... promoting to array')
value = [value]
self.tcex.log.debug(f'... {name} = {value!r}')
self.outlist.append((name, value, output_type))
def lookup_dns(self) -> None:
"""Run the App main logic.
This method should contain the core logic of the App.
"""
questions = tc_argcheck(
self.tcex.rargs, 'questions', label='Question(s)', required=True, tcex=self.tcex
)
if not isinstance(questions, list):
questions = [questions]
record_types = tc_argcheck(self.tcex.rargs, 'record_types', required=True, tcex=self.tcex)
if not isinstance(record_types, list):
record_types = [record_types]
record_types = [x for x in record_types if x]
if not record_types:
self.fail('At least one resource record type is required.')
self.tcex.log.debug(f'Questions: {questions!r}, rrtypes {record_types!r}')
for question in questions:
if isinstance(question, dict): # must be tcentity
entity_type = question.get('type')
question = question.get('value')
if entity_type == 'EmailAddress':
if '@' not in question:
self.tcex.log.warning(f'Invalid EmailAddress {question} -- Skipped')
continue
question = question.split('@', 1)[1]
elif entity_type == 'Address':
pass
elif entity_type == 'Host':
pass
elif entity_type.upper() == 'URL':
question = urlparse(question).netloc
else:
self.tcex.log.warning(f'Unexpected indicator type {entity_type} -- Skipped')
continue
for rrtype in record_types:
self.questions.append((question, rrtype))
self.tcex.log.debug(f'Queuing {len(self.questions)} for resolution')
self.batch_resolve()
result = {}
cnames = {}
valid_questions = set()
invalid_questions = set()
for answer in self.answers:
question, cname, answers = answer
qname, rrtype = question
rrdict = result.get(qname, {})
result[qname] = rrdict
alist = rrdict.get(rrtype, [])
rrdict[rrtype] = alist
if answers:
valid_questions.add(qname)
for a in answers:
if a not in alist:
alist.append(a)
if qname not in cnames:
cnames[qname] = cname
for answer in self.answers:
question, cname, answers = answer
qname, rrtype = question
if qname not in valid_questions:
invalid_questions.add(qname)
self.add_output('dns.result.json', result, jsonify=True)
self.add_output('dns.valid', sorted(list(valid_questions)))
self.add_output('dns.invalid', sorted(list(invalid_questions)))
def fail_on_no_results(self):
"""Return True if fail_on_no_results is set"""
return tc_argcheck(
self.tcex.args, 'fail_on_no_results', types=bool, default=False, tcex=self.tcex
)