From 27dd411b2753eed2e38c44378156e533b625901e Mon Sep 17 00:00:00 2001 From: Brandon Miller Date: Sat, 28 Nov 2020 00:41:10 -0500 Subject: [PATCH] Cleaned up code with flake8 linter --- __init__.py | 4 ++-- binja/binja_export.py | 45 +++++++++++++++++-------------------- binja/binja_import.py | 52 ++++++++++++++++++++++++++----------------- ida/ida_export.py | 41 +++++++++++++++++++++------------- ida/ida_import.py | 26 ++++++++++++++-------- 5 files changed, 96 insertions(+), 72 deletions(-) diff --git a/__init__.py b/__init__.py index 10927ed..c640f30 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,5 @@ -from binaryninja import * -from .binja import * +from binaryninja import PluginCommand +from .binja import import_data_in_background, export_data_in_background PluginCommand.register( 'bnida: Import analysis data', diff --git a/binja/binja_export.py b/binja/binja_export.py index 1ab82b1..4b55920 100644 --- a/binja/binja_export.py +++ b/binja/binja_export.py @@ -1,15 +1,10 @@ -import json -from optparse import OptionParser -from binaryninja import * -from collections import OrderedDict - """ Exports analysis data from a BN database to a bnida JSON file """ -__author__ = 'zznop' -__copyright__ = 'Copyright 2018, zznop0x90@gmail.com' -__license__ = 'MIT' +import json +from binaryninja import SaveFileNameField, get_form_input, BackgroundTaskThread +from collections import OrderedDict class GetOptions(object): @@ -21,14 +16,15 @@ def __init__(self): else: self.json_file = json_file.result + class ExportInBackground(BackgroundTaskThread): def __init__(self, bv, options): global task BackgroundTaskThread.__init__(self, 'Exporting data from BN', False) self.json_file = options.json_file - self.options = options - self.bv = bv - task = self + self.options = options + self.bv = bv + task = self def get_sections(self): """ @@ -41,8 +37,8 @@ def get_sections(self): for section_name in self.bv.sections: section = self.bv.get_section_by_name(section_name) sections[section.name] = { - 'start' : section.start, - 'end' : section.end + 'start': section.start, + 'end': section.end } return sections @@ -80,7 +76,6 @@ def get_function_comments(self): comments = {} for func in self.bv: - current_function = {} if func.comment: comments[func.start] = func.comment @@ -125,8 +120,8 @@ def get_structures(self): for member in typ.structure.members: members[member.name] = {} members[member.name]['offset'] = member.offset - members[member.name]['size'] = member.type.width - members[member.name]['type'] = '' + members[member.name]['size'] = member.type.width + members[member.name]['type'] = '' for token in member.type.tokens: members[member.name]['type'] += str(token) @@ -141,20 +136,22 @@ def run(self): Export analysis data to bnida JSON file """ - print('[*] Exporting analysis data to {}'.format(self.options.json_file)) - json_array = {} - json_array['sections'] = self.get_sections() - json_array['names'] = self.get_names() - json_array['functions'] = self.get_functions() - json_array['func_comments'] = self.get_function_comments() - json_array['line_comments'] = self.get_line_comments() - json_array['structs'] = self.get_structures() + print('[*] Exporting analysis data to {}'.format( + self.options.json_file)) + json_array = {} + json_array['sections'] = self.get_sections() + json_array['names'] = self.get_names() + json_array['functions'] = self.get_functions() + json_array['func_comments'] = self.get_function_comments() + json_array['line_comments'] = self.get_line_comments() + json_array['structs'] = self.get_structures() with open(self.options.json_file, 'w+') as f: json.dump(json_array, f, indent=4) print('[+] Done exporting analysis data') + def export_data_in_background(bv): """ Export data in background from BN UI diff --git a/binja/binja_import.py b/binja/binja_import.py index 214b085..6647b1c 100644 --- a/binja/binja_import.py +++ b/binja/binja_import.py @@ -1,14 +1,12 @@ -import json -from collections import OrderedDict -from binaryninja import * - """ Imports analysis data from a bnida json file into a Binary Ninja database """ -__author__ = 'zznop' -__copyright__ = 'Copyright 2018, zznop0x90@gmail.com' -__license__ = 'MIT' +import json +from collections import OrderedDict +from binaryninja import (OpenFileNameField, get_form_input, + BackgroundTaskThread, types, Type, SymbolType, Symbol, + Architecture) class GetOptions(object): @@ -32,7 +30,8 @@ class ImportInBackground(BackgroundTaskThread): def __init__(self, bv, options): global task - BackgroundTaskThread.__init__(self, 'Importing data from bnida JSON file', False) + BackgroundTaskThread.__init__( + self, 'Importing data from bnida JSON file', False) self.json_file = options.json_file self.bv = bv self.options = options @@ -56,13 +55,14 @@ def adjust_addr(self, sections, addr): # Make sure the section was found (this check should always pass) if section_name is None: - print('Section not found in analysis data for addr: {:08x}'.format(addr)) + print('Section not found in analysis data: {:08x}'.format(addr)) return None # Retrieve section start in BN bn_section = self.bv.get_section_by_name(section_name) if bn_section is None: - print('Section not found in BN - name:{} addr:{:08x}'.format(section_name, addr)) + print('Section not found in BN - name:{} addr:{:08x}'.format( + section_name, addr)) return None # Adjust if needed @@ -116,9 +116,11 @@ def import_function_comments(self, comments, sections): def import_line_comments(self, comments, sections): """ - Import line comments into BN database. Binary Ninja uses BinaryView.set_function_at or Function.set_comment_at - for line comments in data sections or line comments in Functions, respectively. Therefore, we check if the - addr is contained in a function and use the appropriate API. + Import line comments into BN database. Binary Ninja uses + BinaryView.set_function_at or Function.set_comment_at for line + comments in data sections or line comments in Functions, respectively. + Therefore, we check if the addr is contained in a function and use the + appropriate API. :param comments: Dict containing line comments :param sections: Dict containing section info @@ -148,11 +150,13 @@ def import_structures(self, structs): struct = types.Structure() for member_name, member_info in struct_info['members'].items(): try: - typ, _ = self.bv.parse_type_string('{}'.format(member_info['type'])) + typ, _ = self.bv.parse_type_string('{}'.format( + member_info['type'])) except SyntaxError: - print('Failed to apply type ({}) to member ({}) in structure ({})'.format( + print('Failed to apply type ({}) to member ({}): ({})'.format( member_info['type'], member_name, struct_name)) - typ, _ = self.bv.parse_type_string('uint8_t [{}]'.format(member_info['size'])) + typ, _ = self.bv.parse_type_string('uint8_t [{}]'.format( + member_info['size'])) struct.insert(int(member_info['offset']), typ, member_name) self.bv.define_user_type(struct_name, Type.structure_type(struct)) @@ -171,9 +175,11 @@ def import_names(self, names, sections): continue if self.bv.get_function_at(addr): - self.bv.define_user_symbol(Symbol(SymbolType.FunctionSymbol, addr, name)) + self.bv.define_user_symbol( + Symbol(SymbolType.FunctionSymbol, addr, name)) else: - self.bv.define_user_symbol(Symbol(SymbolType.DataSymbol, addr, name)) + self.bv.define_user_symbol( + Symbol(SymbolType.DataSymbol, addr, name)) def get_architectures(self): """ @@ -193,20 +199,24 @@ def run(self): Open JSON file and apply analysis data to BN database """ - print('[*] Importing analysis data from {}'.format(self.options.json_file)) + print('[*] Importing analysis data from {}'.format( + self.options.json_file)) json_array = self.open_json_file(self.options.json_file) if self.bv.platform is None: print('[!] Platform has not been set, cannot import analysis data') return self.import_functions(json_array['functions'], json_array['sections']) - self.import_function_comments(json_array['func_comments'], json_array['sections']) - self.import_line_comments(json_array['line_comments'], json_array['sections']) + self.import_function_comments( + json_array['func_comments'], json_array['sections']) + self.import_line_comments( + json_array['line_comments'], json_array['sections']) self.import_names(json_array['names'], json_array['sections']) self.import_structures(json_array['structs']) self.bv.update_analysis_and_wait() print('[+] Done importing analysis data') + def import_data_in_background(bv): """ Registered plugin command handler diff --git a/ida/ida_export.py b/ida/ida_export.py index 7f35b7f..3aac56a 100644 --- a/ida/ida_export.py +++ b/ida/ida_export.py @@ -1,19 +1,17 @@ +""" +Exports analysis data from IDA to a bnida json file +""" + import ida_struct import ida_kernwin import ida_segment import ida_bytes import idautils +import ida_funcs +import idaapi import json from collections import OrderedDict -""" -Exports analysis data from IDA to a bnida json file -""" - -__author__ = 'zznop' -__copyright__ = 'Copyright 2018, zznop0x90@gmail.com' -__license__ = 'MIT' - def get_single_comment(regular, repeatable): """ @@ -31,6 +29,7 @@ def get_single_comment(regular, repeatable): elif regular is not None and repeatable is None: return regular + def get_single_function_comment(ea): """ Get function comment @@ -44,6 +43,7 @@ def get_single_function_comment(ea): repeatable = ida_funcs.get_func_cmt(func, True) return get_single_comment(regular, repeatable) + def get_single_line_comment(ea): """ Get line comment @@ -57,6 +57,7 @@ def get_single_line_comment(ea): cmt = get_single_comment(regular, repeatable) return cmt + def get_function_comments(): """ Get function comments from IDA database @@ -72,6 +73,7 @@ def get_function_comments(): return comments + def get_functions(): """ Get function start addresses @@ -85,6 +87,7 @@ def get_functions(): return func_addrs + def get_line_comments(): """ Iterate through every address in a segment and check for comments @@ -122,6 +125,7 @@ def get_names(): return symbols + def get_sections(): """ Get section names and start/end addrs from IDA database @@ -137,12 +141,13 @@ def get_sections(): continue curr = {} - curr['start'] = segm.start_ea - curr['end'] = segm.end_ea + curr['start'] = segm.start_ea + curr['end'] = segm.end_ea sections[name] = curr return sections + def get_member_type(struct, idx): """ Retrieve the type information for the struct member @@ -180,6 +185,7 @@ def get_member_type(struct, idx): return typ + def get_struct_members(struct, sid): """ Get members belonging to a structure by structure ID @@ -199,10 +205,11 @@ def get_struct_members(struct, sid): # Type isn't set so make it a byte array members[name]['type'] = 'uint8_t [{}]'.format(size) members[name]['offset'] = offset - members[name]['size'] = size + members[name]['size'] = size return members + def get_structs(): """ Get structures from IDA database @@ -214,11 +221,12 @@ def get_structs(): for idx, sid, name in idautils.Structs(): struct = ida_struct.get_struc(sid) structs[name] = {} - structs[name]['size'] = ida_struct.get_struc_size(struct) + structs[name]['size'] = ida_struct.get_struc_size(struct) structs[name]['members'] = get_struct_members(struct, sid) return structs + def main(json_file): """ Construct a json file containing analysis data from an IDA database @@ -228,16 +236,17 @@ def main(json_file): json_array = {} print('[*] Exporting analysis data to {}'.format(json_file)) - json_array['sections'] = get_sections() - json_array['functions'] = get_functions() + json_array['sections'] = get_sections() + json_array['functions'] = get_functions() json_array['func_comments'] = get_function_comments() json_array['line_comments'] = get_line_comments() - json_array['names'] = get_names() - json_array['structs'] = get_structs() + json_array['names'] = get_names() + json_array['structs'] = get_structs() print('[+] Done exporting analysis data') with open(json_file, 'w') as f: f.write(json.dumps(json_array, indent=4)) + if __name__ == '__main__': main(ida_kernwin.ask_file(1, '*.json', 'Export file name')) diff --git a/ida/ida_import.py b/ida/ida_import.py index f16b126..2c292d5 100644 --- a/ida/ida_import.py +++ b/ida/ida_import.py @@ -1,3 +1,7 @@ +""" +Imports analysis data from a bnida JSON file to IDA +""" + import idc import idautils import ida_kernwin @@ -5,15 +9,9 @@ import ida_segment import ida_struct import ida_bytes +import ida_funcs import json -""" -Imports analysis data from a bnida JSON file to IDA -""" - -__author__ = 'zznop' -__copyright__ = 'Copyright 2018, zznop0x90@gmail.com' -__license__ = 'MIT' def get_flag_from_type(typ): """ @@ -31,6 +29,7 @@ def get_flag_from_type(typ): else: return ida_bytes.byte_flag() + def sanitize_name(name): """ Remove characters from names that IDA doesn't like @@ -43,6 +42,7 @@ def sanitize_name(name): name = name.replace('@', '_') return name + def adjust_addr(sections, addr): """ Adjust the address if there are differences in section base addresses @@ -75,6 +75,7 @@ def adjust_addr(sections, addr): print('Section not found - name:{} addr:{:08x}'.format(section_name, addr)) return None + def import_functions(functions, sections): """ Create functions from bnida analysis data @@ -91,9 +92,10 @@ def import_functions(functions, sections): if ida_funcs.get_func(addr): continue - if ida_funcs.add_func(addr) != True: + if not ida_funcs.add_func(addr): print('Failed to create function at offset:{:08x}'.format(addr)) + def import_function_comments(comments, sections): """ Import function comments into IDA @@ -114,6 +116,7 @@ def import_function_comments(comments, sections): ida_funcs.set_func_cmt(func, comment, False) + def import_line_comments(comments, sections): """ Import line comments @@ -126,6 +129,7 @@ def import_line_comments(comments, sections): addr = adjust_addr(sections, int(addr)) ida_bytes.set_cmt(addr, comment, 0) + def import_names(names, sections): """ Import symbol names @@ -143,6 +147,7 @@ def import_names(names, sections): if idc.get_name_ea_simple(name) == idaapi.BADADDR: idc.set_name(addr, name) + def import_structures(structures): """ Import structures @@ -186,6 +191,7 @@ def get_json(json_file): print('Failed to parse json file {} {}'.format(json_file, e)) return json_array + def main(json_file): """ Import analysis data from bnida JSON file @@ -198,11 +204,13 @@ def main(json_file): print('[*] Importing analysis data from {}'.format(json_file)) import_functions(json_array['functions'], json_array['sections']) - import_function_comments(json_array['func_comments'], json_array['sections']) + import_function_comments( + json_array['func_comments'], json_array['sections']) import_line_comments(json_array['line_comments'], json_array['sections']) import_names(json_array['names'], json_array['sections']) import_structures(json_array['structs']) print('[+] Done importing analysis data') + if __name__ == '__main__': main(ida_kernwin.ask_file(1, '*.json', 'Import file name'))