autoformat with black

1 files changed, 50 insertions, 34 deletions

diff --git a/unicorn_mode/helper_scripts/ida_context_loader.py b/unicorn_mode/helper_scripts/ida_context_loader.py
index 31d47a90..d7984c77 100644
--- a/unicorn_mode/helper_scripts/ida_context_loader.py
+++ b/unicorn_mode/helper_scripts/ida_context_loader.py
@@ -34,13 +34,11 @@ import ida_segment
 
 
 class ContextLoaderError(Exception):
-    """Base "catch all" exception for this script
-    """
+    """Base "catch all" exception for this script"""
 
 
 class ArchNotSupportedError(ContextLoaderError):
-    """Exception raised if the input file CPU architecture isn't supported fully
-    """
+    """Exception raised if the input file CPU architecture isn't supported fully"""
 
 
 def parse_mapping_index(filepath: str):
@@ -51,13 +49,16 @@ def parse_mapping_index(filepath: str):
     """
 
     if filepath is None:
-        raise ContextLoaderError('_index.json file was not selected')
+        raise ContextLoaderError("_index.json file was not selected")
 
     try:
-        with open(filepath, 'rb') as _file:
+        with open(filepath, "rb") as _file:
             return json.load(_file)
     except Exception as ex:
-        raise ContextLoaderError('Failed to parse json file {}'.format(filepath)) from ex
+        raise ContextLoaderError(
+            "Failed to parse json file {}".format(filepath)
+        ) from ex
+
 
 def get_input_name():
     """Get the name of the input file
@@ -68,19 +69,21 @@ def get_input_name():
     input_filepath = ida_nalt.get_input_file_path()
     return Path(input_filepath).name
 
+
 def write_segment_bytes(start: int, filepath: str):
-    """"Read data from context file and write it to the IDA segment
+    """ "Read data from context file and write it to the IDA segment
 
     :param start: Start address
     :param filepath: Path to context file
     """
 
-    with open(filepath, 'rb') as _file:
+    with open(filepath, "rb") as _file:
         data = _file.read()
 
     decompressed_data = zlib.decompress(data)
     ida_bytes.put_bytes(start, decompressed_data)
 
+
 def create_segment(context_dir: str, segment: dict, is_be: bool):
     """Create segment in IDA and map in the data from the file
 
@@ -90,23 +93,30 @@ def create_segment(context_dir: str, segment: dict, is_be: bool):
     """
 
     input_name = get_input_name()
-    if Path(segment['name']).name != input_name:
+    if Path(segment["name"]).name != input_name:
         ida_seg = idaapi.segment_t()
-        ida_seg.start_ea = segment['start']
-        ida_seg.end_ea = segment['end']
+        ida_seg.start_ea = segment["start"]
+        ida_seg.end_ea = segment["end"]
         ida_seg.bitness = 1 if is_be else 0
-        if segment['permissions']['r']:
+        if segment["permissions"]["r"]:
             ida_seg.perm |= ida_segment.SEGPERM_READ
-        if segment['permissions']['w']:
+        if segment["permissions"]["w"]:
             ida_seg.perm |= ida_segment.SEGPERM_WRITE
-        if segment['permissions']['x']:
+        if segment["permissions"]["x"]:
             ida_seg.perm |= ida_segment.SEGPERM_EXEC
-            idaapi.add_segm_ex(ida_seg, Path(segment['name']).name, 'CODE', idaapi.ADDSEG_OR_DIE)
+            idaapi.add_segm_ex(
+                ida_seg, Path(segment["name"]).name, "CODE", idaapi.ADDSEG_OR_DIE
+            )
         else:
-            idaapi.add_segm_ex(ida_seg, Path(segment['name']).name, 'DATA', idaapi.ADDSEG_OR_DIE)
+            idaapi.add_segm_ex(
+                ida_seg, Path(segment["name"]).name, "DATA", idaapi.ADDSEG_OR_DIE
+            )
+
+    if segment["content_file"]:
+        write_segment_bytes(
+            segment["start"], PurePath(context_dir, segment["content_file"])
+        )
 
-    if segment['content_file']:
-        write_segment_bytes(segment['start'], PurePath(context_dir, segment['content_file']))
 
 def create_segments(index: dict, context_dir: str):
     """Iterate segments in index JSON, create the segment in IDA, and map in the data from the file
@@ -117,9 +127,10 @@ def create_segments(index: dict, context_dir: str):
 
     info = idaapi.get_inf_structure()
     is_be = info.is_be()
-    for segment in index['segments']:
+    for segment in index["segments"]:
         create_segment(context_dir, segment, is_be)
 
+
 def rebase_program(index: dict):
     """Rebase the program to the offset specified in the context _index.json
 
@@ -128,20 +139,21 @@ def rebase_program(index: dict):
 
     input_name = get_input_name()
     new_base = None
-    for segment in index['segments']:
-        if not segment['name']:
+    for segment in index["segments"]:
+        if not segment["name"]:
             continue
 
-        segment_name = Path(segment['name']).name
+        segment_name = Path(segment["name"]).name
         if input_name == segment_name:
-            new_base = segment['start']
+            new_base = segment["start"]
             break
 
     if not new_base:
-        raise ContextLoaderError('Input file is not in _index.json')
+        raise ContextLoaderError("Input file is not in _index.json")
 
     current_base = idaapi.get_imagebase()
-    ida_segment.rebase_program(new_base-current_base, 8)
+    ida_segment.rebase_program(new_base - current_base, 8)
+
 
 def get_pc_by_arch(index: dict) -> int:
     """Queries the input file CPU architecture and attempts to lookup the address of the program
@@ -153,13 +165,14 @@ def get_pc_by_arch(index: dict) -> int:
 
     progctr = None
     info = idaapi.get_inf_structure()
-    if info.procname == 'metapc':
+    if info.procname == "metapc":
         if info.is_64bit():
-            progctr = index['regs']['rax']
+            progctr = index["regs"]["rax"]
         elif info.is_32bit():
-            progctr = index['regs']['eax']
+            progctr = index["regs"]["eax"]
     return progctr
 
+
 def write_reg_info(index: dict):
     """Write register info as line comment at instruction pointed to by the program counter and
     change focus to that location
@@ -167,17 +180,19 @@ def write_reg_info(index: dict):
     :param index: _index.json JSON data
     """
 
-    cmt = ''
-    for reg, val in index['regs'].items():
+    cmt = ""
+    for reg, val in index["regs"].items():
         cmt += f"{reg.ljust(6)} : {hex(val)}\n"
 
     progctr = get_pc_by_arch(index)
     if progctr is None:
         raise ArchNotSupportedError(
-            'Architecture not fully supported, skipping register status comment')
+            "Architecture not fully supported, skipping register status comment"
+        )
     ida_bytes.set_cmt(progctr, cmt, 0)
     ida_kernwin.jumpto(progctr)
 
+
 def main(filepath):
     """Main - parse _index.json input and map context files into the database
 
@@ -193,5 +208,6 @@ def main(filepath):
     except ContextLoaderError as ex:
         print(ex)
 
-if __name__ == '__main__':
-    main(ida_kernwin.ask_file(1, '*.json', 'Import file name'))
+
+if __name__ == "__main__":
+    main(ida_kernwin.ask_file(1, "*.json", "Import file name"))