Commit e663f5ff authored by Štěpán Balážik's avatar Štěpán Balážik Committed by Stepan Balazik

scenario: clean-up raw data handling

parent ea966ba2
......@@ -260,8 +260,7 @@ class Entry:
self.fired = 0
# RAW
self.raw_data = None # type: Optional[bytes]
self.is_raw_data_entry = self.process_raw()
self.raw_data = self.process_raw()
# MATCH
self.match_fields = self.process_match()
......@@ -286,10 +285,9 @@ class Entry:
def process_raw(self):
try:
self.raw_data = binascii.unhexlify(self.node["/raw"].value)
return True
return binascii.unhexlify(self.node["/raw"].value)
except KeyError:
return False
return None
def process_match(self):
try:
......@@ -374,7 +372,7 @@ class Entry:
def __str__(self):
txt = 'ENTRY_BEGIN\n'
if not self.is_raw_data_entry:
if self.raw_data is None:
txt += 'MATCH {0}\n'.format(' '.join(self.match_fields))
txt += 'ADJUST {0}\n'.format(' '.join(self.adjust_fields))
txt += 'REPLY {rcode} {flags}\n'.format(
......@@ -390,7 +388,7 @@ class Entry:
for rr in sect:
txt += str(rr)
txt += '\n'
if self.is_raw_data_entry:
if self.raw_data is not None:
txt += 'RAW\n'
if self.raw_data:
txt += binascii.hexlify(self.raw_data)
......@@ -459,7 +457,7 @@ class Entry:
raise ValueError("%s, \"%s\": %s" % (self.node.span, code, errstr))
def cmp_raw(self, raw_value):
assert self.is_raw_data_entry
assert self.raw_data is not None
expected = None
if self.raw_data is not None:
expected = binascii.hexlify(self.raw_data)
......@@ -472,7 +470,7 @@ class Entry:
def reply(self, query) -> Optional[DNSBlob]:
if 'do_not_answer' in self.adjust_fields:
return None
if self.is_raw_data_entry:
if self.raw_data is not None:
copy_id = 'raw_data' in self.adjust_fields
assert self.raw_data is not None
return DNSReplyRaw(self.raw_data, query, copy_id)
......@@ -681,7 +679,7 @@ class Step:
if not self.data:
raise ValueError("response definition required")
expected = self.data[0]
if expected.is_raw_data_entry is True:
if expected.raw_data is not None:
self.log.debug("raw answer: %s", ctx.last_raw_answer.to_text())
expected.cmp_raw(ctx.last_raw_answer)
else:
......@@ -721,7 +719,7 @@ class Step:
"""
if not self.data:
raise ValueError("query definition required")
if self.data[0].is_raw_data_entry is True:
if self.data[0].raw_data is not None:
data_to_wire = self.data[0].raw_data
else:
# Don't use a message copy as the EDNS data portion is not copied.
......@@ -754,7 +752,7 @@ class Step:
time.sleep(0.1)
# Wait for a response for a reasonable time
answer = None
if not self.data[0].is_raw_data_entry:
if self.data[0].raw_data is None:
while True:
if (datetime.now() - tstart).total_seconds() > 5:
raise RuntimeError("Server took too long to respond")
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment