Skip to content

Commit

Permalink
Made the metadata writing code less prone to accidentally lose metadata.
Browse files Browse the repository at this point in the history
  • Loading branch information
Thomas Dullien committed Nov 12, 2018
1 parent 1f4611b commit ac8be7f
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 15 deletions.
10 changes: 6 additions & 4 deletions pybindings/binary_ninja_plugin/modules/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,9 @@
]

# Default location of simhash database
default_sim_hash_location = '/tmp/example.simhash'
default_sim_hash_location = '/var/tmp/simhash.index'
default_sim_hash_meta_location = default_sim_hash_location+'.meta'


class Plugin:
def __init__(self):
self.sim_hash_location = None
Expand Down Expand Up @@ -100,7 +99,7 @@ def get_exec_id(self, filename):
self.exec_id_cache[filename] = long(h.hexdigest()[0:16], 16)
return self.exec_id_cache[filename]

def save_single_function_hash(self, bv, search_index, function):
def save_single_function_hash(self, bv, search_index, function, write_meta=True):
"""
Save the hash of a given function into a given search index.
"""
Expand All @@ -111,6 +110,8 @@ def save_single_function_hash(self, bv, search_index, function):
search_index.add_function(h1, h2, exec_id, function.start)
bn.log_info('[+] Added function <{:x}:0x{:x} {:x}-{:x}> to search index.'.format(exec_id, function.start, h1, h2))
self.metadata.add(exec_id, function.start, bv.file.filename, function.name)
if write_meta:
self.metadata.__save__()
else:
bn.log_info('[-] Did not add function <{:x}:0x{:x}> to search index.'.format(exec_id, function.start))

Expand Down Expand Up @@ -144,7 +145,8 @@ def save_all_functions(self, bv, current_function):
"""
search_index = self.init_index(bv, current_function)
for function in bv.functions:
self.save_single_function_hash(bv, search_index, function)
self.save_single_function_hash(bv, search_index, function, False)
self.metadata.__save__()

def add_report_from_result(self, results, report, address, minimal_match = 100):
results = [ r for r in results if r[0] > minimal_match ]
Expand Down
52 changes: 41 additions & 11 deletions pybindings/binary_ninja_plugin/modules/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,20 +9,28 @@

class Metadata:
def __init__(self, meta_location):
self.data = {}
self.data = self.read_metadata_dict(meta_location)
self.sim_hash_meta_location = meta_location

if os.path.isfile(self.sim_hash_meta_location):
def read_metadata_dict(self, meta_location):
data = {}
if os.path.isfile(meta_location):
with open(self.sim_hash_meta_location, "r") as fh:
for line in fh.xreadlines():
tokens = line.split()
if tokens:
self.add(long(tokens[0], 16), long(tokens[2], 16), base64.b64decode(tokens[3]), tokens[1])
try:
decoded_token = base64.b64decode(tokens[3])
except:
bn.log_info("[-] Failure to b64decode '%s'!" % tokens[3])
decoded_token = tokens[3]
data[(exe_id, address)] = (file_name, function_name)
bn.log_info("[+] Loaded metadata for %d functions." % len(data.keys()))
return data

def add(self, exe_id, address, file_name, function_name):
bn.log_info("[*] Adding metadata. exe_id:{:x}, address:{:x}, file_name:{}, function_name:{}".format(exe_id, address, file_name, function_name))
self.data[(exe_id, address)] = (file_name, function_name)
self.__save__()

def get(self, exe_id, address):
if self.data.has_key((exe_id, address)):
Expand All @@ -31,11 +39,33 @@ def get(self, exe_id, address):
return None

def __save__(self):
"""
When saving, we should be conservative -- we do not want to delete existing
metadata, instead we want to add any extra metadata we have collected.
"""
exists = os.path.isfile(self.sim_hash_meta_location)
if exists:
data = self.read_metadata_dict(self.sim_hash_meta_location)
else:
data = {}
for key, value in self.data.iteritems():
data[key] = value
# Now write the file again.
with open(self.sim_hash_meta_location, "w+") as fh:
for key, value in self.data.iteritems():
fh.write("{exe_id:x} {function_name} {address:x} {file_name}\n".format(
exe_id = key[0],
function_name =value[1],
address = key[1],
file_name = base64.b64encode(value[0])
))
for key, value in data.iteritems():
try:
fh.write("{exe_id:x} {function_name} {address:x} {file_name}\n".format(
exe_id = key[0],
function_name =value[1],
address = key[1],
file_name = base64.b64encode(value[0])
))
except:
# Keep on writing even if there was an exception.
pass

def __del__(self):
"""
Save all metadata prior to exiting.
"""
self.__save__()

0 comments on commit ac8be7f

Please sign in to comment.