diff options
Diffstat (limited to 'unicorn_mode')
| -rw-r--r-- | unicorn_mode/helper_scripts/ida_context_loader.py | 84 |
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")) |