HackRF-Treasure-Chest/Software/Universal Radio Hacker/tests/awre/test_length_engine.py
2022-09-22 13:46:47 -07:00

168 lines
6.8 KiB
Python

import random
from tests.awre.AWRETestCase import AWRETestCase
from urh.awre.FormatFinder import FormatFinder
from urh.awre.MessageTypeBuilder import MessageTypeBuilder
from urh.awre.ProtocolGenerator import ProtocolGenerator
from urh.awre.engines.LengthEngine import LengthEngine
from urh.signalprocessing.FieldType import FieldType
from urh.signalprocessing.ProtocoLabel import ProtocolLabel
class TestLengthEngine(AWRETestCase):
def test_simple_protocol(self):
"""
Test a simple protocol with
preamble, sync and length field (8 bit) and some random data
:return:
"""
mb = MessageTypeBuilder("simple_length_test")
mb.add_label(FieldType.Function.PREAMBLE, 8)
mb.add_label(FieldType.Function.SYNC, 16)
mb.add_label(FieldType.Function.LENGTH, 8)
num_messages_by_data_length = {8: 5, 16: 10, 32: 15}
pg = ProtocolGenerator([mb.message_type],
syncs_by_mt={mb.message_type: "0x9a9d"})
random.seed(0)
for data_length, num_messages in num_messages_by_data_length.items():
for i in range(num_messages):
pg.generate_message(data="".join([random.choice(["0", "1"]) for _ in range(data_length)]))
#self.save_protocol("simple_length", pg)
self.clear_message_types(pg.protocol.messages)
ff = FormatFinder(pg.protocol.messages)
length_engine = LengthEngine(ff.bitvectors)
highscored_ranges = length_engine.find(n_gram_length=8)
self.assertEqual(len(highscored_ranges), 3)
ff.perform_iteration()
self.assertEqual(len(ff.message_types), 1)
self.assertGreater(len(ff.message_types[0]), 0)
label = ff.message_types[0].get_first_label_with_type(FieldType.Function.LENGTH)
self.assertEqual(label.start, 24)
self.assertEqual(label.length, 8)
def test_easy_protocol(self):
"""
preamble, sync, sequence number, length field (8 bit) and some random data
:return:
"""
mb = MessageTypeBuilder("easy_length_test")
mb.add_label(FieldType.Function.PREAMBLE, 16)
mb.add_label(FieldType.Function.SYNC, 16)
mb.add_label(FieldType.Function.LENGTH, 8)
mb.add_label(FieldType.Function.SEQUENCE_NUMBER, 8)
num_messages_by_data_length = {32: 10, 64: 15, 16: 5, 24: 7}
pg = ProtocolGenerator([mb.message_type],
preambles_by_mt={mb.message_type: "10" * 8},
syncs_by_mt={mb.message_type: "0xcafe"})
for data_length, num_messages in num_messages_by_data_length.items():
for i in range(num_messages):
if i % 4 == 0:
data = "1" * data_length
elif i % 4 == 1:
data = "0" * data_length
elif i % 4 == 2:
data = "10" * (data_length // 2)
else:
data = "01" * (data_length // 2)
pg.generate_message(data=data)
#self.save_protocol("easy_length", pg)
self.clear_message_types(pg.protocol.messages)
ff = FormatFinder(pg.protocol.messages)
length_engine = LengthEngine(ff.bitvectors)
highscored_ranges = length_engine.find(n_gram_length=8)
self.assertEqual(len(highscored_ranges), 4)
ff.perform_iteration()
self.assertEqual(len(ff.message_types), 1)
self.assertGreater(len(ff.message_types[0]), 0)
label = ff.message_types[0].get_first_label_with_type(FieldType.Function.LENGTH)
self.assertIsInstance(label, ProtocolLabel)
self.assertEqual(label.start, 32)
self.assertEqual(label.length, 8)
def test_medium_protocol(self):
"""
Protocol with two message types. Length field only present in one of them
:return:
"""
mb1 = MessageTypeBuilder("data")
mb1.add_label(FieldType.Function.PREAMBLE, 8)
mb1.add_label(FieldType.Function.SYNC, 8)
mb1.add_label(FieldType.Function.LENGTH, 8)
mb1.add_label(FieldType.Function.SEQUENCE_NUMBER, 8)
mb2 = MessageTypeBuilder("ack")
mb2.add_label(FieldType.Function.PREAMBLE, 8)
mb2.add_label(FieldType.Function.SYNC, 8)
pg = ProtocolGenerator([mb1.message_type, mb2.message_type],
syncs_by_mt={mb1.message_type: "11110011",
mb2.message_type: "11110011"})
num_messages_by_data_length = {8: 5, 16: 10, 32: 5}
for data_length, num_messages in num_messages_by_data_length.items():
for i in range(num_messages):
pg.generate_message(data=pg.decimal_to_bits(10 * i, data_length), message_type=mb1.message_type)
pg.generate_message(message_type=mb2.message_type, data="0xaf")
#self.save_protocol("medium_length", pg)
self.clear_message_types(pg.protocol.messages)
ff = FormatFinder(pg.protocol.messages)
ff.perform_iteration()
self.assertEqual(len(ff.message_types), 2)
length_mt = next(
mt for mt in ff.message_types if mt.get_first_label_with_type(FieldType.Function.LENGTH) is not None)
length_label = length_mt.get_first_label_with_type(FieldType.Function.LENGTH)
for i, sync_end in enumerate(ff.sync_ends):
self.assertEqual(sync_end, 16, msg=str(i))
self.assertEqual(16, length_label.start)
self.assertEqual(8, length_label.length)
def test_little_endian_16_bit(self):
mb = MessageTypeBuilder("little_endian_16_length_test")
mb.add_label(FieldType.Function.PREAMBLE, 8)
mb.add_label(FieldType.Function.SYNC, 16)
mb.add_label(FieldType.Function.LENGTH, 16)
num_messages_by_data_length = {256*8: 5, 16: 4, 512: 2}
pg = ProtocolGenerator([mb.message_type],
syncs_by_mt={mb.message_type: "0x9a9d"},
little_endian=True)
random.seed(0)
for data_length, num_messages in num_messages_by_data_length.items():
for i in range(num_messages):
pg.generate_message(data="".join([random.choice(["0", "1"]) for _ in range(data_length)]))
#self.save_protocol("little_endian_16_length_test", pg)
self.clear_message_types(pg.protocol.messages)
ff = FormatFinder(pg.protocol.messages)
length_engine = LengthEngine(ff.bitvectors)
highscored_ranges = length_engine.find(n_gram_length=8)
self.assertEqual(len(highscored_ranges), 3)
ff.perform_iteration()
self.assertEqual(len(ff.message_types), 1)
self.assertGreater(len(ff.message_types[0]), 0)
label = ff.message_types[0].get_first_label_with_type(FieldType.Function.LENGTH)
self.assertEqual(label.start, 24)
self.assertEqual(label.length, 16)