Cleaned up code with flake8 linter

__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',

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, [email protected]'
-__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

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, [email protected]'
-__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

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, [email protected]'
-__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'))