def mock_program(platform=MOCK_PLATFORM,
*,
segments=None,
types=None,
objects=None):
def mock_find_type(kind, name, filename):
if filename:
return None
for type in types:
if type.kind == kind:
try:
type_name = type.name
except AttributeError:
try:
type_name = type.tag
except AttributeError:
continue
if type_name == name:
return type
return None
def mock_object_find(prog, name, flags, filename):
if filename:
return None
for obj in objects:
if obj.name == name:
if obj.value is not None:
if flags & FindObjectFlags.CONSTANT:
break
elif obj.type.kind == TypeKind.FUNCTION:
if flags & FindObjectFlags.FUNCTION:
break
elif flags & FindObjectFlags.VARIABLE:
break
else:
return None
return Object(prog, obj.type, address=obj.address, value=obj.value)
prog = Program(platform)
if segments is not None:
for segment in segments:
if segment.virt_addr is not None:
prog.add_memory_segment(
segment.virt_addr,
len(segment.buf),
functools.partial(mock_memory_read, segment.buf),
)
if segment.phys_addr is not None:
prog.add_memory_segment(
segment.phys_addr,
len(segment.buf),
functools.partial(mock_memory_read, segment.buf),
True,
)
if types is not None:
prog.add_type_finder(mock_find_type)
if objects is not None:
prog.add_object_finder(mock_object_find)
return prog
def test_overlap_segment_head(self):
# Existing segment: |_______|
# New segment: |_______|
prog = Program()
segment1 = unittest.mock.Mock(side_effect=zero_memory_read)
segment2 = unittest.mock.Mock(side_effect=zero_memory_read)
prog.add_memory_segment(0xffff0040, 128, segment1)
prog.add_memory_segment(0xffff0000, 128, segment2)
prog.read(0xffff0000, 192)
segment1.assert_called_once_with(0xffff0080, 64, 64, False)
segment2.assert_called_once_with(0xffff0000, 128, 0, False)
def test_overlap_same_address_larger_size(self):
# Existing segment: |___|
# New segment: |_______|
prog = Program()
segment1 = unittest.mock.Mock(side_effect=zero_memory_read)
segment2 = unittest.mock.Mock(side_effect=zero_memory_read)
prog.add_memory_segment(0xffff0000, 64, segment1)
prog.add_memory_segment(0xffff0000, 128, segment2)
prog.read(0xffff0000, 128)
segment1.assert_not_called()
segment2.assert_called_once_with(0xffff0000, 128, 0, False)
def test_overlap_same_segment(self):
# Existing segment: |_______|
# New segment: |_______|
prog = Program(MOCK_PLATFORM)
segment1 = unittest.mock.Mock(side_effect=zero_memory_read)
segment2 = unittest.mock.Mock(side_effect=zero_memory_read)
prog.add_memory_segment(0xFFFF0000, 128, segment1)
prog.add_memory_segment(0xFFFF0000, 128, segment2)
prog.read(0xFFFF0000, 128)
segment1.assert_not_called()
segment2.assert_called_once_with(0xFFFF0000, 128, 0, False)
def mock_program(arch=MOCK_ARCH, *, segments=None, types=None, symbols=None):
def mock_find_type(kind, name, filename):
if filename:
return None
for type in types:
if type.kind == kind:
try:
type_name = type.name
except AttributeError:
try:
type_name = type.tag
except AttributeError:
continue
if type_name == name:
return type
return None
def mock_symbol_find(name, flags, filename):
if filename:
return None
for sym_name, sym in symbols:
if sym_name == name:
if sym.value is not None or sym.is_enumerator:
if flags & FindObjectFlags.CONSTANT:
break
elif sym.type.kind == TypeKind.FUNCTION:
if flags & FindObjectFlags.FUNCTION:
break
elif flags & FindObjectFlags.VARIABLE:
break
else:
return None
return sym
prog = Program(arch)
if segments is not None:
for segment in segments:
if segment.virt_addr is not None:
prog.add_memory_segment(
segment.virt_addr, len(segment.buf),
functools.partial(mock_memory_read, segment.buf))
if segment.phys_addr is not None:
prog.add_memory_segment(
segment.phys_addr, len(segment.buf),
functools.partial(mock_memory_read, segment.buf), True)
if types is not None:
prog.add_type_finder(mock_find_type)
if symbols is not None:
prog.add_symbol_finder(mock_symbol_find)
return prog
def test_overlap_within_segment(self):
# Existing segment: |_______|
# New segment: |___|
prog = Program()
segment1 = unittest.mock.Mock(side_effect=zero_memory_read)
segment2 = unittest.mock.Mock(side_effect=zero_memory_read)
prog.add_memory_segment(0xffff0000, 128, segment1)
prog.add_memory_segment(0xffff0020, 64, segment2)
prog.read(0xffff0000, 128)
segment1.assert_has_calls([
unittest.mock.call(0xffff0000, 32, 00, False),
unittest.mock.call(0xffff0060, 32, 96, False),
])
segment2.assert_called_once_with(0xffff0020, 64, 0, False)
def test_overlap_segment_head_and_tail(self):
# Existing segment: |_______||_______|
# New segment: |_______|
prog = Program(MOCK_PLATFORM)
segment1 = unittest.mock.Mock(side_effect=zero_memory_read)
segment2 = unittest.mock.Mock(side_effect=zero_memory_read)
segment3 = unittest.mock.Mock(side_effect=zero_memory_read)
prog.add_memory_segment(0xFFFF0000, 128, segment1)
prog.add_memory_segment(0xFFFF0080, 128, segment2)
prog.add_memory_segment(0xFFFF0040, 128, segment3)
prog.read(0xFFFF0000, 256)
segment1.assert_called_once_with(0xFFFF0000, 64, 0, False)
segment2.assert_called_once_with(0xFFFF00C0, 64, 64, False)
segment3.assert_called_once_with(0xFFFF0040, 128, 0, False)
def test_overlap_subsume_at_and_after(self):
# Existing segments: |_|_|_|_|
# New segment: |_______|
prog = Program()
segment1 = unittest.mock.Mock(side_effect=zero_memory_read)
segment2 = unittest.mock.Mock(side_effect=zero_memory_read)
prog.add_memory_segment(0xffff0000, 32, segment1)
prog.add_memory_segment(0xffff0020, 32, segment1)
prog.add_memory_segment(0xffff0040, 32, segment1)
prog.add_memory_segment(0xffff0060, 32, segment1)
prog.add_memory_segment(0xffff0000, 128, segment2)
prog.read(0xffff0000, 128)
segment1.assert_not_called()
segment2.assert_called_once_with(0xffff0000, 128, 0, False)
def test_overlap_subsume_after(self):
# Existing segments: |_|_|_|_|
# New segment: |_______|
prog = Program(MOCK_PLATFORM)
segment1 = unittest.mock.Mock(side_effect=zero_memory_read)
segment2 = unittest.mock.Mock(side_effect=zero_memory_read)
segment3 = unittest.mock.Mock(side_effect=zero_memory_read)
prog.add_memory_segment(0xFFFF0020, 32, segment1)
prog.add_memory_segment(0xFFFF0040, 32, segment1)
prog.add_memory_segment(0xFFFF0060, 32, segment1)
prog.add_memory_segment(0xFFFF0080, 64, segment2)
prog.add_memory_segment(0xFFFF0000, 128, segment3)
prog.read(0xFFFF0000, 192)
segment1.assert_not_called()
segment2.assert_called_once_with(0xFFFF0080, 64, 0, False)
segment3.assert_called_once_with(0xFFFF0000, 128, 0, False)
