assert loc_db.get_offset_location(0x1144) == loc_key4_bis
loc_db.consistency_check()
# Names manipulation
loc_key5 = loc_db.add_location()
name1 = "name1"
name2 = "name2"
name3 = "name3"
assert len(loc_db.get_location_names(loc_key5)) == 0
loc_db.add_location_name(loc_key5, name1)
loc_db.add_location_name(loc_key5, name2)
assert name1 in loc_db.names
assert name2 in loc_db.names
assert name1 in loc_db.get_location_names(loc_key5)
assert name2 in loc_db.get_location_names(loc_key5)
assert loc_db.get_name_location(name1) == loc_key5
loc_db.remove_location_name(loc_key5, name1)
assert name1 not in loc_db.names
assert name1 not in loc_db.get_location_names(loc_key5)
try:
loc_db.remove_location_name(loc_key5, name1)
has_raised = False
except KeyError:
has_raised = True
try:
loc_db.add_location_name(loc_key1, name2)
has_raised = False
except KeyError:
has_raised = True
try:
loc_db.add_location(name=name2)
assert loc_db.get_offset_location(0x1144) == loc_key4_bis
loc_db.consistency_check()
# Names manipulation
loc_key5 = loc_db.add_location()
name1 = "name1"
name2 = "name2"
name3 = "name3"
assert len(loc_db.get_location_names(loc_key5)) == 0
loc_db.add_location_name(loc_key5, name1)
loc_db.add_location_name(loc_key5, name2)
assert name1 in loc_db.names
assert name2 in loc_db.names
assert name1 in loc_db.get_location_names(loc_key5)
assert name2 in loc_db.get_location_names(loc_key5)
assert loc_db.get_name_location(name1) == loc_key5
loc_db.remove_location_name(loc_key5, name1)
assert name1 not in loc_db.names
assert name1 not in loc_db.get_location_names(loc_key5)
try:
loc_db.remove_location_name(loc_key5, name1)
has_raised = False
except KeyError:
has_raised = True
try:
loc_db.add_location_name(loc_key1, name2)
has_raised = False
except KeyError:
has_raised = True
try:
loc_db.add_location(name=name2)
addr_main = 0
virt = st
output = st
# Get and parse the source code
with open(args.source) as fstream:
source = fstream.read()
loc_db = LocationDB()
asmcfg, loc_db = parse_asm.parse_txt(machine.mn, attrib, source, loc_db)
# Fix shellcode addrs
loc_db.set_location_offset(loc_db.get_name_location("main"), addr_main)
if args.PE:
loc_db.set_location_offset(loc_db.get_or_create_name_location("MessageBoxA"),
pe.DirImport.get_funcvirt('USER32.dll',
'MessageBoxA'))
# Print and graph firsts blocks before patching it
for block in asmcfg.blocks:
print block
open("graph.dot", "w").write(asmcfg.dot())
# Apply patches
patches = asmblock.asm_resolve_final(machine.mn,
asmcfg,
loc_db,
