class MessageTypeBuilder(object):
def __init__(self, name: str):
self.name = name
self.message_type = MessageType(name)
def add_label(self,
label_type: FieldType.Function,
length: int,
name: str = None):
try:
start = self.message_type[-1].end
color_index = self.message_type[-1].color_index + 1
except IndexError:
start, color_index = 0, 0
if name is None:
name = label_type.value
lbl = ProtocolLabel(name,
start,
start + length - 1,
color_index,
field_type=FieldType(label_type.name, label_type))
self.message_type.append(lbl)
def add_checksum_label(self,
length,
checksum,
data_start=None,
data_end=None,
name: str = None):
label_type = FieldType.Function.CHECKSUM
try:
start = self.message_type[-1].end
color_index = self.message_type[-1].color_index + 1
except IndexError:
start, color_index = 0, 0
if name is None:
name = label_type.value
if data_start is None:
# End of sync or preamble
sync_label = self.message_type.get_first_label_with_type(
FieldType.Function.SYNC)
if sync_label:
data_start = sync_label.end
else:
preamble_label = self.message_type.get_first_label_with_type(
FieldType.Function.PREAMBLE)
if preamble_label:
data_start = preamble_label.end
else:
data_start = 0
if data_end is None:
data_end = start
lbl = ChecksumLabel(name,
start,
start + length - 1,
color_index,
field_type=FieldType(label_type.name, label_type))
lbl.data_ranges = [(data_start, data_end)]
lbl.checksum = checksum
self.message_type.append(lbl)
def perform_iteration_for_message_type(self, message_type: MessageType):
"""
Perform a field inference iteration for messages of the given message type
This routine will return newly found fields as a set of Common Ranges
:param message_type:
:rtype: set of CommonRange
"""
indices = self.existing_message_types[message_type]
engines = []
# We can take an arbitrary sync end to correct the already labeled fields for this message type,
# because if the existing labels would have different sync positions,
# they would not belong to the same message type in the first place
sync_end = self.sync_ends[indices[0]] if indices else 0
already_labeled = [(lbl.start - sync_end, lbl.end - sync_end)
for lbl in message_type if lbl.start >= sync_end]
if not message_type.get_first_label_with_type(
FieldType.Function.LENGTH):
engines.append(
LengthEngine([self.bitvectors[i] for i in indices],
already_labeled=already_labeled))
if not message_type.get_first_label_with_type(
FieldType.Function.SRC_ADDRESS):
engines.append(
AddressEngine([self.hexvectors[i] for i in indices],
[self.participant_indices[i] for i in indices],
self.known_participant_addresses,
already_labeled=already_labeled))
elif not message_type.get_first_label_with_type(
FieldType.Function.DST_ADDRESS):
engines.append(
AddressEngine([self.hexvectors[i] for i in indices],
[self.participant_indices[i] for i in indices],
self.known_participant_addresses,
already_labeled=already_labeled,
src_field_present=True))
if not message_type.get_first_label_with_type(
FieldType.Function.SEQUENCE_NUMBER):
engines.append(
SequenceNumberEngine([self.bitvectors[i] for i in indices],
already_labeled=already_labeled))
if not message_type.get_first_label_with_type(
FieldType.Function.CHECKSUM):
# If checksum was not found in first iteration, it will also not be found in next one
if self.current_iteration == 0:
engines.append(
ChecksumEngine([self.bitvectors[i] for i in indices],
already_labeled=already_labeled))
result = set()
for engine in engines:
high_scored_ranges = engine.find() # type: list[CommonRange]
high_scored_ranges = self.retransform_message_indices(
high_scored_ranges, indices, self.sync_ends)
merged_ranges = self.merge_common_ranges(high_scored_ranges)
result.update(merged_ranges)
return result
