From c68dc3bf029e47f003777da8734423a51e267c4e Mon Sep 17 00:00:00 2001 From: Michael Hunhoff Date: Mon, 6 Jul 2020 21:01:26 -0600 Subject: [PATCH] IDA code maintenance --- capa/features/extractors/ida/__init__.py | 47 +++--- capa/features/extractors/ida/basicblock.py | 78 ++++------ capa/features/extractors/ida/file.py | 37 ++--- capa/features/extractors/ida/function.py | 40 ++---- capa/features/extractors/ida/helpers.py | 157 ++++++++++++--------- capa/features/extractors/ida/insn.py | 129 ++++++----------- capa/ida/explorer/item.py | 10 +- capa/ida/explorer/model.py | 45 +++--- capa/ida/explorer/proxy.py | 1 + capa/ida/explorer/view.py | 47 ++---- capa/ida/ida_capa_explorer.py | 78 +++++----- 11 files changed, 312 insertions(+), 357 deletions(-) diff --git a/capa/features/extractors/ida/__init__.py b/capa/features/extractors/ida/__init__.py index 5091fb34..b9ea7086 100644 --- a/capa/features/extractors/ida/__init__.py +++ b/capa/features/extractors/ida/__init__.py @@ -8,34 +8,30 @@ import capa.features.extractors.ida.insn import capa.features.extractors.ida.helpers import capa.features.extractors.ida.function import capa.features.extractors.ida.basicblock + from capa.features.extractors import FeatureExtractor -def get_va(self): - if isinstance(self, idaapi.BasicBlock): +def get_ea(self): + """ """ + if isinstance(self, (idaapi.BasicBlock, idaapi.func_t)): return self.start_ea - - if isinstance(self, idaapi.func_t): - return self.start_ea - if isinstance(self, idaapi.insn_t): return self.ea - raise TypeError -def add_va_int_cast(o): +def add_ea_int_cast(o): """ dynamically add a cast-to-int (`__int__`) method to the given object - that returns the value of the `.va` property. + that returns the value of the `.ea` property. this bit of skullduggery lets use cast viv-utils objects as ints. the correct way of doing this is to update viv-utils (or subclass the objects here). """ - - if sys.version_info >= (3, 0): - setattr(o, "__int__", types.MethodType(get_va, o)) + if sys.version_info[0] >= 3: + setattr(o, "__int__", types.MethodType(get_ea, o)) else: - setattr(o, "__int__", types.MethodType(get_va, o, type(o))) + setattr(o, "__int__", types.MethodType(get_ea, o, type(o))) return o @@ -47,29 +43,30 @@ class IdaFeatureExtractor(FeatureExtractor): return idaapi.get_imagebase() def extract_file_features(self): - for feature, va in capa.features.extractors.ida.file.extract_features(): - yield feature, va + for (feature, ea) in capa.features.extractors.ida.file.extract_features(): + yield feature, ea def get_functions(self): - for f in capa.features.extractors.ida.helpers.get_functions(ignore_thunks=True, ignore_libs=True): - yield add_va_int_cast(f) + # ignore library functions and thunk functions as identified by IDA + for f in capa.features.extractors.ida.helpers.get_functions(skip_thunks=True, skip_libs=True): + yield add_ea_int_cast(f) def extract_function_features(self, f): - for feature, va in capa.features.extractors.ida.function.extract_features(f): - yield feature, va + for (feature, ea) in capa.features.extractors.ida.function.extract_features(f): + yield feature, ea def get_basic_blocks(self, f): for bb in idaapi.FlowChart(f, flags=idaapi.FC_PREDS): - yield add_va_int_cast(bb) + yield add_ea_int_cast(bb) def extract_basic_block_features(self, f, bb): - for feature, va in capa.features.extractors.ida.basicblock.extract_features(f, bb): - yield feature, va + for (feature, ea) in capa.features.extractors.ida.basicblock.extract_features(f, bb): + yield feature, ea def get_instructions(self, f, bb): for insn in capa.features.extractors.ida.helpers.get_instructions_in_range(bb.start_ea, bb.end_ea): - yield add_va_int_cast(insn) + yield add_ea_int_cast(insn) def extract_insn_features(self, f, bb, insn): - for feature, va in capa.features.extractors.ida.insn.extract_features(f, bb, insn): - yield feature, va + for (feature, ea) in capa.features.extractors.ida.insn.extract_features(f, bb, insn): + yield feature, ea diff --git a/capa/features/extractors/ida/basicblock.py b/capa/features/extractors/ida/basicblock.py index b935318a..75520b60 100644 --- a/capa/features/extractors/ida/basicblock.py +++ b/capa/features/extractors/ida/basicblock.py @@ -1,11 +1,10 @@ import sys -import pprint import string import struct -import idc import idaapi -import idautils + +import capa.features.extractors.ida.helpers from capa.features import Characteristic from capa.features.basicblock import BasicBlock @@ -13,13 +12,14 @@ from capa.features.extractors.ida import helpers from capa.features.extractors.helpers import MIN_STACKSTRING_LEN -def _ida_get_printable_len(op): +def get_printable_len(op): """ Return string length if all operand bytes are ascii or utf16-le printable args: op (IDA op_t) """ - op_val = helpers.mask_op_val(op) + op_val = capa.features.extractors.ida.helpers.mask_op_val(op) + if op.dtype == idaapi.dt_byte: chars = struct.pack("= (3, 0): + def is_printable_ascii(chars): + if sys.version_info[0] >= 3: return all(c < 127 and chr(c) in string.printable for c in chars) else: return all(ord(c) < 127 and c in string.printable for c in chars) - def _is_printable_utf16le(chars): - if sys.version_info >= (3, 0): + def is_printable_utf16le(chars): + if sys.version_info[0] >= 3: if all(c == 0x00 for c in chars[1::2]): - return _is_printable_ascii(chars[::2]) + return is_printable_ascii(chars[::2]) else: if all(c == "\x00" for c in chars[1::2]): - return _is_printable_ascii(chars[::2]) + return is_printable_ascii(chars[::2]) - if _is_printable_ascii(chars): + if is_printable_ascii(chars): return idaapi.get_dtype_size(op.dtype) - if _is_printable_utf16le(chars): - return idaapi.get_dtype_size(op.dtype) / 2 + if is_printable_utf16le(chars): + return idaapi.get_dtype_size(op.dtype) // 2 return 0 -def _is_mov_imm_to_stack(insn): +def is_mov_imm_to_stack(insn): """ verify instruction moves immediate onto stack args: @@ -72,8 +72,7 @@ def _is_mov_imm_to_stack(insn): return True - -def _ida_bb_contains_stackstring(f, bb): +def bb_contains_stackstring(f, bb): """ check basic block for stackstring indicators true if basic block contains enough moves of constant bytes to the stack @@ -83,14 +82,11 @@ def _ida_bb_contains_stackstring(f, bb): bb (IDA BasicBlock) """ count = 0 - - for insn in helpers.get_instructions_in_range(bb.start_ea, bb.end_ea): - if _is_mov_imm_to_stack(insn): - count += _ida_get_printable_len(insn.Op2) - + for insn in capa.features.extractors.ida.helpers.get_instructions_in_range(bb.start_ea, bb.end_ea): + if is_mov_imm_to_stack(insn): + count += get_printable_len(insn.Op2) if count > MIN_STACKSTRING_LEN: return True - return False @@ -101,29 +97,10 @@ def extract_bb_stackstring(f, bb): f (IDA func_t) bb (IDA BasicBlock) """ - if _ida_bb_contains_stackstring(f, bb): + if bb_contains_stackstring(f, bb): yield Characteristic("stack string"), bb.start_ea -def _ida_bb_contains_tight_loop(f, bb): - """ check basic block for stackstring indicators - - true if last instruction in basic block branches to basic block start - - args: - f (IDA func_t) - bb (IDA BasicBlock) - """ - bb_end = idc.prev_head(bb.end_ea) - - if bb.start_ea < bb_end: - for ref in idautils.CodeRefsFrom(bb_end, True): - if ref == bb.start_ea: - return True - - return False - - def extract_bb_tight_loop(f, bb): """ extract tight loop indicators from a basic block @@ -131,7 +108,7 @@ def extract_bb_tight_loop(f, bb): f (IDA func_t) bb (IDA BasicBlock) """ - if _ida_bb_contains_tight_loop(f, bb): + if capa.features.extractors.ida.helpers.is_basic_block_tight_loop(bb): yield Characteristic("tight loop"), bb.start_ea @@ -142,11 +119,10 @@ def extract_features(f, bb): f (IDA func_t) bb (IDA BasicBlock) """ - yield BasicBlock(), bb.start_ea - for bb_handler in BASIC_BLOCK_HANDLERS: - for feature, va in bb_handler(f, bb): - yield feature, va + for (feature, ea) in bb_handler(f, bb): + yield feature, ea + yield BasicBlock(), bb.start_ea BASIC_BLOCK_HANDLERS = ( @@ -157,11 +133,11 @@ BASIC_BLOCK_HANDLERS = ( def main(): features = [] - - for f in helpers.get_functions(ignore_thunks=True, ignore_libs=True): + for f in helpers.get_functions(skip_thunks=True, skip_libs=True): for bb in idaapi.FlowChart(f, flags=idaapi.FC_PREDS): features.extend(list(extract_features(f, bb))) - + + import pprint pprint.pprint(features) diff --git a/capa/features/extractors/ida/file.py b/capa/features/extractors/ida/file.py index 4824d770..c8f0bf1e 100644 --- a/capa/features/extractors/ida/file.py +++ b/capa/features/extractors/ida/file.py @@ -1,4 +1,3 @@ -import pprint import struct import idc @@ -8,11 +7,12 @@ import idautils import capa.features.extractors.helpers import capa.features.extractors.strings import capa.features.extractors.ida.helpers + from capa.features import String, Characteristic from capa.features.file import Export, Import, Section -def _ida_check_segment_for_pe(seg): +def check_segment_for_pe(seg): """ check segment for embedded PE adapted for IDA from: @@ -66,18 +66,14 @@ def extract_file_embedded_pe(): - '-R' from console - Check 'Load resource sections' when opening binary in IDA manually """ - for seg in capa.features.extractors.ida.helpers.get_segments(): - if seg.is_header_segm(): - # IDA may load header segments, skip if present - continue - - for ea, _ in _ida_check_segment_for_pe(seg): + for seg in capa.features.extractors.ida.helpers.get_segments(skip_header_segments=True): + for (ea, _) in check_segment_for_pe(seg): yield Characteristic("embedded pe"), ea def extract_file_export_names(): """ extract function exports """ - for _, _, ea, name in idautils.Entries(): + for (_, _, ea, name) in idautils.Entries(): yield Export(name), ea @@ -92,15 +88,12 @@ def extract_file_import_names(): - modulename.importname - importname """ - for ea, imp_info in capa.features.extractors.ida.helpers.get_file_imports().items(): - dllname, name, ordi = imp_info - - if name: - yield Import("%s.%s" % (dllname, name)), ea - yield Import(name), ea - - if ordi: - yield Import("%s.#%s" % (dllname, str(ordi))), ea + for (ea, info) in capa.features.extractors.ida.helpers.get_file_imports().items(): + if info[1]: + yield Import("%s.%s" % (info[0], info[1])), ea + yield Import(info[1]), ea + if info[2]: + yield Import("%s.#%s" % (info[0], str(info[2]))), ea def extract_file_section_names(): @@ -110,11 +103,7 @@ def extract_file_section_names(): - '-R' from console - Check 'Load resource sections' when opening binary in IDA manually """ - for seg in capa.features.extractors.ida.helpers.get_segments(): - if seg.is_header_segm(): - # IDA may load header segments, skip if present - continue - + for seg in capa.features.extractors.ida.helpers.get_segments(skip_header_segments=True): yield Section(idaapi.get_segm_name(seg)), seg.start_ea @@ -152,6 +141,8 @@ FILE_HANDLERS = ( def main(): + """ """ + import pprint pprint.pprint(list(extract_features())) diff --git a/capa/features/extractors/ida/function.py b/capa/features/extractors/ida/function.py index 416aa495..72d29b44 100644 --- a/capa/features/extractors/ida/function.py +++ b/capa/features/extractors/ida/function.py @@ -1,34 +1,19 @@ import idaapi import idautils +import capa.features.extractors.ida.helpers + from capa.features import Characteristic from capa.features.extractors import loops -def _ida_function_contains_switch(f): - """ check a function for switch statement indicators - - adapted from: - https://reverseengineering.stackexchange.com/questions/17548/calc-switch-cases-in-idapython-cant-iterate-over-results?rq=1 - - arg: - f (IDA func_t) - """ - for start, end in idautils.Chunks(f.start_ea): - for head in idautils.Heads(start, end): - if idaapi.get_switch_info(head): - return True - - return False - - def extract_function_switch(f): """ extract switch indicators from a function arg: f (IDA func_t) """ - if _ida_function_contains_switch(f): + if capa.features.extractors.ida.helpers.is_function_switch_statement(f): yield Characteristic("switch"), f.start_ea @@ -49,10 +34,12 @@ def extract_function_loop(f): f (IDA func_t) """ edges = [] + + # construct control flow graph for bb in idaapi.FlowChart(f): map(lambda s: edges.append((bb.start_ea, s.start_ea)), bb.succs()) - if edges and loops.has_loop(edges): + if loops.has_loop(edges): yield Characteristic("loop"), f.start_ea @@ -62,10 +49,8 @@ def extract_recursive_call(f): args: f (IDA func_t) """ - for ref in idautils.CodeRefsTo(f.start_ea, True): - if f.contains(ref): - yield Characteristic("recursive call"), f.start_ea - break + if capa.features.extractors.ida.helpers.is_function_recursive(f): + yield Characteristic("recursive call"), f.start_ea def extract_features(f): @@ -75,19 +60,20 @@ def extract_features(f): f (IDA func_t) """ for func_handler in FUNCTION_HANDLERS: - for feature, va in func_handler(f): - yield feature, va + for (feature, ea) in func_handler(f): + yield feature, ea FUNCTION_HANDLERS = (extract_function_calls_to, extract_function_switch, extract_function_loop, extract_recursive_call) def main(): + """ """ features = [] - - for f in helpers.get_functions(ignore_thunks=True, ignore_libs=True): + for f in capa.features.extractors.ida.get_functions(skip_thunks=True, skip_libs=True): features.extend(list(extract_features(f))) + import pprint pprint.pprint(features) diff --git a/capa/features/extractors/ida/helpers.py b/capa/features/extractors/ida/helpers.py index 0ddd450d..a28c9a63 100644 --- a/capa/features/extractors/ida/helpers.py +++ b/capa/features/extractors/ida/helpers.py @@ -14,13 +14,13 @@ def find_byte_sequence(start, end, seq): end: max virtual address seq: bytes to search e.g. b'\x01\x03' """ - if sys.version_info >= (3, 0): + if sys.version_info[0] >= 3: return idaapi.find_binary(start, end, " ".join(["%02x" % b for b in seq]), 0, idaapi.SEARCH_DOWN) else: return idaapi.find_binary(start, end, " ".join(["%02x" % ord(b) for b in seq]), 0, idaapi.SEARCH_DOWN) -def get_functions(start=None, end=None, ignore_thunks=False, ignore_libs=False): +def get_functions(start=None, end=None, skip_thunks=False, skip_libs=False): """ get functions, range optional args: @@ -32,21 +32,19 @@ def get_functions(start=None, end=None, ignore_thunks=False, ignore_libs=False): """ for ea in idautils.Functions(start=start, end=end): f = idaapi.get_func(ea) - - if ignore_thunks and f.flags & idaapi.FUNC_THUNK: - continue - - if ignore_libs and f.flags & idaapi.FUNC_LIB: - continue - - yield f + if not (skip_thunks and (f.flags & idaapi.FUNC_THUNK) or skip_libs and (f.flags & idaapi.FUNC_LIB)): + yield f -def get_segments(): - """ Get list of segments (sections) in the binary image """ +def get_segments(skip_header_segments=False): + """ get list of segments (sections) in the binary image + + args: + skip_header_segments: IDA may load header segments - skip if set + """ for n in range(idaapi.get_segm_qty()): seg = idaapi.getnseg(n) - if seg: + if seg and not (skip_header_segments and seg.is_header_segm()): yield seg @@ -70,24 +68,24 @@ def get_segment_buffer(seg): def get_file_imports(): """ get file imports """ - _imports = {} + imports = {} for idx in range(idaapi.get_import_module_qty()): - dllname = idaapi.get_import_module_name(idx) + library = idaapi.get_import_module_name(idx) - if not dllname: + if not library: continue - def _inspect_import(ea, name, ordi): - if name and name.startswith("__imp_"): + def inspect_import(ea, function, ordinal): + if function and function.startswith("__imp_"): # handle mangled names starting - name = name[len("__imp_") :] - _imports[ea] = (dllname.lower(), name, ordi) + function = function[len("__imp_") :] + imports[ea] = (library.lower(), function, ordinal) return True - idaapi.enum_import_names(idx, _inspect_import) + idaapi.enum_import_names(idx, inspect_import) - return _imports + return imports def get_instructions_in_range(start, end): @@ -100,9 +98,9 @@ def get_instructions_in_range(start, end): (insn_t*) """ for head in idautils.Heads(start, end): - inst = idautils.DecodeInstruction(head) - if inst: - yield inst + insn = idautils.DecodeInstruction(head) + if insn: + yield insn def is_operand_equal(op1, op2): @@ -133,7 +131,16 @@ def is_operand_equal(op1, op2): def is_basic_block_equal(bb1, bb2): """ compare two IDA BasicBlock """ - return bb1.start_ea == bb2.start_ea and bb1.end_ea == bb2.end_ea and bb1.type == bb2.type + if bb1.start_ea != bb2.start_ea: + return False + + if bb1.end_ea != bb2.end_ea: + return False + + if bb1.type != bb2.type: + return False + + return True def basic_block_size(bb): @@ -142,6 +149,7 @@ def basic_block_size(bb): def read_bytes_at(ea, count): + """ """ segm_end = idc.get_segm_end(ea) if ea + count > segm_end: return idc.get_bytes(ea, segm_end - ea) @@ -163,7 +171,7 @@ def find_string_at(ea, min=4): return found except UnicodeDecodeError: pass - return None + return "" def get_op_phrase_info(op): @@ -173,7 +181,7 @@ def get_op_phrase_info(op): https://github.com/tmr232/Sark/blob/master/sark/code/instruction.py#L28-L73 """ if op.type not in (idaapi.o_phrase, idaapi.o_displ): - return + return {} scale = 1 << ((op.specflag2 & 0xC0) >> 6) offset = op.addr @@ -191,7 +199,7 @@ def get_op_phrase_info(op): if index & 4: index += 8 else: - return + return {} if (index == base == idautils.procregs.sp.reg) and (scale == 1): # HACK: This is a really ugly hack. For some reason, phrases of the form `[esp + ...]` (`sp`, `rsp` as well) @@ -215,25 +223,19 @@ def is_op_read(insn, op): def is_sp_modified(insn): """ determine if instruction modifies SP, ESP, RSP """ - for op in get_insn_ops(insn, op_type=(idaapi.o_reg,)): - if op.reg != idautils.procregs.sp.reg: - continue - - if is_op_write(insn, op): + for op in get_insn_ops(insn, target_ops=(idaapi.o_reg,)): + if op.reg == idautils.procregs.sp.reg and is_op_write(insn, op): + # register is stack and written return True - return False def is_bp_modified(insn): """ check if instruction modifies BP, EBP, RBP """ - for op in get_insn_ops(insn, op_type=(idaapi.o_reg,)): - if op.reg != idautils.procregs.bp.reg: - continue - - if is_op_write(insn, op): + for op in get_insn_ops(insn, target_ops=(idaapi.o_reg,)): + if op.reg == idautils.procregs.bp.reg and is_op_write(insn, op): + # register is base and written return True - return False @@ -242,33 +244,26 @@ def is_frame_register(reg): return reg in (idautils.procregs.sp.reg, idautils.procregs.bp.reg) -def get_insn_ops(insn, op_type=None): +def get_insn_ops(insn, target_ops=()): """ yield op_t for instruction, filter on type if specified """ for op in insn.ops: if op.type == idaapi.o_void: # avoid looping all 6 ops if only subset exists break - - if op_type and op.type not in op_type: + if target_ops and op.type not in target_ops: continue - yield op -def ea_flags(ea): - """ retrieve processor flags for a given address """ - return idaapi.get_flags(ea) - - -def is_op_stack_var(ea, n): +def is_op_stack_var(ea, index): """ check if operand is a stack variable """ - return idaapi.is_stkvar(ea_flags(ea), n) + return idaapi.is_stkvar(idaapi.get_flags(ea), index) def mask_op_val(op): - """ mask off a value based on data type + """ mask value by data type - necesssary due to a bug in 64-bit + necessary due to a bug in AMD64 Example: .rsrc:0054C12C mov [ebp+var_4], 0FFFFFFFFh @@ -282,15 +277,49 @@ def mask_op_val(op): idaapi.dt_dword: 0xFFFFFFFF, idaapi.dt_qword: 0xFFFFFFFFFFFFFFFF, } - - mask = masks.get(op.dtype, None) - - if not mask: - raise ValueError("No support for operand data type 0x%x" % op.dtype) - - return mask & op.value + return masks.get(op.dtype, op.value) & op.value -def ea_to_offset(ea): - """ convert virtual address to file offset """ - return idaapi.get_fileregion_offset(ea) +def is_function_recursive(f): + """ check if function is recursive + + args: + f (IDA func_t) + """ + for ref in idautils.CodeRefsTo(f.start_ea, True): + if f.contains(ref): + return True + return False + + +def is_function_switch_statement(f): + """ check a function for switch statement indicators + + adapted from: + https://reverseengineering.stackexchange.com/questions/17548/calc-switch-cases-in-idapython-cant-iterate-over-results?rq=1 + + arg: + f (IDA func_t) + """ + for (start, end) in idautils.Chunks(f.start_ea): + for head in idautils.Heads(start, end): + if idaapi.get_switch_info(head): + return True + return False + + +def is_basic_block_tight_loop(bb): + """ check basic block loops to self + + true if last instruction in basic block branches to basic block start + + args: + f (IDA func_t) + bb (IDA BasicBlock) + """ + bb_end = idc.prev_head(bb.end_ea) + if bb.start_ea < bb_end: + for ref in idautils.CodeRefsFrom(bb_end, True): + if ref == bb.start_ea: + return True + return False diff --git a/capa/features/extractors/ida/insn.py b/capa/features/extractors/ida/insn.py index 4a784e3c..f42e8b06 100644 --- a/capa/features/extractors/ida/insn.py +++ b/capa/features/extractors/ida/insn.py @@ -1,11 +1,10 @@ -import pprint - import idc import idaapi import idautils import capa.features.extractors.helpers import capa.features.extractors.ida.helpers + from capa.features import MAX_BYTES_FEATURE_SIZE, Bytes, String, Characteristic from capa.features.insn import Number, Offset, Mnemonic @@ -13,34 +12,32 @@ _file_imports_cache = None def get_imports(): + """ """ global _file_imports_cache if _file_imports_cache is None: _file_imports_cache = capa.features.extractors.ida.helpers.get_file_imports() return _file_imports_cache -def _check_for_api_call(insn): +def check_for_api_call(insn): """ check instruction for API call """ if not idaapi.is_call_insn(insn): return - for call_ref in idautils.CodeRefsFrom(insn.ea, False): - imp = get_imports().get(call_ref, None) - - if imp: - yield "%s.%s" % (imp[0], imp[1]) + for ref in idautils.CodeRefsFrom(insn.ea, False): + info = get_imports().get(ref, ()) + if info: + yield "%s.%s" % (info[0], info[1]) else: - f = idaapi.get_func(call_ref) - - if f and f.flags & idaapi.FUNC_THUNK: - # check if call to thunk - # TODO: first instruction might not always be the thunk - for thunk_ref in idautils.DataRefsFrom(call_ref): + f = idaapi.get_func(ref) + # check if call to thunk + # TODO: first instruction might not always be the thunk + if f and (f.flags & idaapi.FUNC_THUNK): + for thunk_ref in idautils.DataRefsFrom(ref): # TODO: always data ref for thunk?? - imp = get_imports().get(thunk_ref, None) - - if imp: - yield "%s.%s" % (imp[0], imp[1]) + info = get_imports().get(thunk_ref, ()) + if info: + yield "%s.%s" % (info[0], info[1]) def extract_insn_api_features(f, bb, insn): @@ -54,9 +51,9 @@ def extract_insn_api_features(f, bb, insn): example: call dword [0x00473038] """ - for api_name in _check_for_api_call(insn): - for feature, va in capa.features.extractors.helpers.generate_api_features(api_name, insn.ea): - yield feature, va + for api in check_for_api_call(insn): + for (feature, ea) in capa.features.extractors.helpers.generate_api_features(api, insn.ea): + yield feature, ea def extract_insn_number_features(f, bb, insn): @@ -80,14 +77,10 @@ def extract_insn_number_features(f, bb, insn): # .text:00401145 add esp, 0Ch return - for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, op_type=(idaapi.o_imm,)): - op_val = capa.features.extractors.ida.helpers.mask_op_val(op) - - if idaapi.is_mapped(op_val): - # assume valid address is not a constant - continue - - yield Number(op_val), insn.ea + for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_imm,)): + const = capa.features.extractors.ida.helpers.mask_op_val(op) + if not idaapi.is_mapped(const): + yield Number(const), insn.ea def extract_insn_bytes_features(f, bb, insn): @@ -107,9 +100,8 @@ def extract_insn_bytes_features(f, bb, insn): for ref in idautils.DataRefsFrom(insn.ea): extracted_bytes = capa.features.extractors.ida.helpers.read_bytes_at(ref, MAX_BYTES_FEATURE_SIZE) - if extracted_bytes: - if not capa.features.extractors.helpers.all_zeros(extracted_bytes): - yield Bytes(extracted_bytes), insn.ea + if extracted_bytes and not capa.features.extractors.helpers.all_zeros(extracted_bytes): + yield Bytes(extracted_bytes), insn.ea def extract_insn_string_features(f, bb, insn): @@ -140,51 +132,36 @@ def extract_insn_offset_features(f, bb, insn): example: .text:0040112F cmp [esi+4], ebx """ - for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, op_type=(idaapi.o_phrase, idaapi.o_displ)): + for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_phrase, idaapi.o_displ)): if capa.features.extractors.ida.helpers.is_op_stack_var(insn.ea, op.n): - # skip stack offsets continue - p_info = capa.features.extractors.ida.helpers.get_op_phrase_info(op) - - if not p_info: - continue - - op_off = p_info["offset"] - + op_off = p_info.get("offset", 0) if 0 == op_off: - # TODO: Do we want to record offset of zero? continue - if idaapi.is_mapped(op_off): # Ignore: # mov esi, dword_1005B148[esi] continue - - # TODO: Do we handle two's complement? yield Offset(op_off), insn.ea -def _contains_stack_cookie_keywords(s): +def contains_stack_cookie_keywords(s): """ check if string contains stack cookie keywords Examples: xor ecx, ebp ; StackCookie - mov eax, ___security_cookie """ if not s: return False - s = s.strip().lower() - if "cookie" not in s: return False - return any(keyword in s for keyword in ("stack", "security")) -def _bb_stack_cookie_registers(bb): +def bb_stack_cookie_registers(bb): """ scan basic block for stack cookie operations yield registers ids that may have been used for stack cookie operations @@ -211,26 +188,25 @@ def _bb_stack_cookie_registers(bb): TODO: this is expensive, but necessary?... """ for insn in capa.features.extractors.ida.helpers.get_instructions_in_range(bb.start_ea, bb.end_ea): - if _contains_stack_cookie_keywords(idc.GetDisasm(insn.ea)): - for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, op_type=(idaapi.o_reg,)): + if contains_stack_cookie_keywords(idc.GetDisasm(insn.ea)): + for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_reg,)): if capa.features.extractors.ida.helpers.is_op_write(insn, op): # only include modified registers yield op.reg -def _is_nzxor_stack_cookie(f, bb, insn): +def is_nzxor_stack_cookie(f, bb, insn): """ check if nzxor is related to stack cookie """ - if _contains_stack_cookie_keywords(idaapi.get_cmt(insn.ea, False)): + if contains_stack_cookie_keywords(idaapi.get_cmt(insn.ea, False)): # Example: # xor ecx, ebp ; StackCookie return True - - if any(op_reg in _bb_stack_cookie_registers(bb) for op_reg in (insn.Op1.reg, insn.Op2.reg)): + stack_cookie_regs = tuple(bb_stack_cookie_registers(bb)) + if any(op_reg in stack_cookie_regs for op_reg in (insn.Op1.reg, insn.Op2.reg)): # Example: # mov eax, ___security_cookie # xor eax, ebp return True - return False @@ -246,13 +222,10 @@ def extract_insn_nzxor_characteristic_features(f, bb, insn): """ if insn.itype != idaapi.NN_xor: return - if capa.features.extractors.ida.helpers.is_operand_equal(insn.Op1, insn.Op2): return - - if _is_nzxor_stack_cookie(f, bb, insn): + if is_nzxor_stack_cookie(f, bb, insn): return - yield Characteristic("nzxor"), insn.ea @@ -278,8 +251,8 @@ def extract_insn_peb_access_characteristic_features(f, bb, insn): if insn.itype not in (idaapi.NN_push, idaapi.NN_mov): return - if all(map(lambda op: op.type != idaapi.o_mem, insn.ops)): - # try to optimize for only memory referencese + if any(map(lambda op: op.type != idaapi.o_mem, insn.ops)): + # try to optimize for only memory references return disasm = idc.GetDisasm(insn.ea) @@ -296,7 +269,7 @@ def extract_insn_segment_access_features(f, bb, insn): IDA should be able to do this... """ if all(map(lambda op: op.type != idaapi.o_mem, insn.ops)): - # try to optimize for only memory referencese + # try to optimize for only memory references return disasm = idc.GetDisasm(insn.ea) @@ -322,14 +295,11 @@ def extract_insn_cross_section_cflow(f, bb, insn): if ref in get_imports().keys(): # ignore API calls continue - if not idaapi.getseg(ref): # handle IDA API bug continue - if idaapi.getseg(ref) == idaapi.getseg(insn.ea): continue - yield Characteristic("cross section flow"), insn.ea @@ -343,12 +313,9 @@ def extract_function_calls_from(f, bb, insn): bb (IDA BasicBlock) insn (IDA insn_t) """ - if not idaapi.is_call_insn(insn): - # ignore jmp, etc. - return - - for ref in idautils.CodeRefsFrom(insn.ea, False): - yield Characteristic("calls from"), ref + if idaapi.is_call_insn(insn): + for ref in idautils.CodeRefsFrom(insn.ea, False): + yield Characteristic("calls from"), ref def extract_function_indirect_call_characteristic_features(f, bb, insn): @@ -363,10 +330,7 @@ def extract_function_indirect_call_characteristic_features(f, bb, insn): bb (IDA BasicBlock) insn (IDA insn_t) """ - if not idaapi.is_call_insn(insn): - return - - if idc.get_operand_type(insn.ea, 0) in (idc.o_reg, idc.o_phrase, idc.o_displ): + if idaapi.is_call_insn(insn) and idc.get_operand_type(insn.ea, 0) in (idc.o_reg, idc.o_phrase, idc.o_displ): yield Characteristic("indirect call"), insn.ea @@ -379,8 +343,8 @@ def extract_features(f, bb, insn): insn (IDA insn_t) """ for inst_handler in INSTRUCTION_HANDLERS: - for feature, va in inst_handler(f, bb, insn): - yield feature, va + for feature, ea in inst_handler(f, bb, insn): + yield feature, ea INSTRUCTION_HANDLERS = ( @@ -400,13 +364,14 @@ INSTRUCTION_HANDLERS = ( def main(): + """ """ features = [] - - for f in capa.features.extractors.ida.helpers.get_functions(ignore_thunks=True, ignore_libs=True): + for f in capa.features.extractors.ida.helpers.get_functions(skip_thunks=True, skip_libs=True): for bb in idaapi.FlowChart(f, flags=idaapi.FC_PREDS): for insn in capa.features.extractors.ida.helpers.get_instructions_in_range(bb.start_ea, bb.end_ea): features.extend(list(extract_features(f, bb, insn))) + import pprint pprint.pprint(features) diff --git a/capa/ida/explorer/item.py b/capa/ida/explorer/item.py index e9df8ded..7e0d5706 100644 --- a/capa/ida/explorer/item.py +++ b/capa/ida/explorer/item.py @@ -190,10 +190,11 @@ class CapaExplorerFunctionItem(CapaExplorerDataItem): class CapaExplorerSubscopeItem(CapaExplorerDataItem): - + """ store data relevant to subscope """ fmt = "subscope(%s)" def __init__(self, parent, scope): + """ """ super(CapaExplorerSubscopeItem, self).__init__(parent, [self.fmt % scope, "", ""]) @@ -220,11 +221,14 @@ class CapaExplorerFeatureItem(CapaExplorerDataItem): """ store data relevant to capa feature result """ def __init__(self, parent, display, location="", details=""): + """ """ location = location_to_hex(location) if location else "" super(CapaExplorerFeatureItem, self).__init__(parent, [display, location, details]) class CapaExplorerInstructionViewItem(CapaExplorerFeatureItem): + """ store data relevant to an instruction preview """ + def __init__(self, parent, display, location): """ """ details = capa.ida.helpers.get_disasm_line(location) @@ -233,6 +237,8 @@ class CapaExplorerInstructionViewItem(CapaExplorerFeatureItem): class CapaExplorerByteViewItem(CapaExplorerFeatureItem): + """ store data relevant to byte preview """ + def __init__(self, parent, display, location): """ """ byte_snap = idaapi.get_bytes(location, 32) @@ -251,6 +257,8 @@ class CapaExplorerByteViewItem(CapaExplorerFeatureItem): class CapaExplorerStringViewItem(CapaExplorerFeatureItem): + """ store data relevant to string preview """ + def __init__(self, parent, display, location): """ """ super(CapaExplorerStringViewItem, self).__init__(parent, display, location=location) diff --git a/capa/ida/explorer/model.py b/capa/ida/explorer/model.py index 4d876fee..e5cb2bff 100644 --- a/capa/ida/explorer/model.py +++ b/capa/ida/explorer/model.py @@ -47,7 +47,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): root_index = self.index(idx, 0, QtCore.QModelIndex()) for model_index in self.iterateChildrenIndexFromRootIndex(root_index, ignore_root=False): model_index.internalPointer().setChecked(False) - self.util_reset_ida_highlighting(model_index.internalPointer(), False) + self.reset_ida_highlighting(model_index.internalPointer(), False) self.dataChanged.emit(model_index, model_index) def clear(self): @@ -239,8 +239,12 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): for idx in range(self.rowCount(child_index)): stack.append(child_index.child(idx, 0)) - def util_reset_ida_highlighting(self, item, checked): - """ """ + def reset_ida_highlighting(self, item, checked): + """ reset IDA highlight for an item + + @param item: capa explorer item + @param checked: indicates item is or not checked + """ if not isinstance( item, (CapaExplorerStringViewItem, CapaExplorerInstructionViewItem, CapaExplorerByteViewItem) ): @@ -281,7 +285,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): # user un/checked box - un/check parent and children for child_index in self.iterateChildrenIndexFromRootIndex(model_index, ignore_root=False): child_index.internalPointer().setChecked(value) - self.util_reset_ida_highlighting(child_index.internalPointer(), value) + self.reset_ida_highlighting(child_index.internalPointer(), value) self.dataChanged.emit(child_index, child_index) return True @@ -428,6 +432,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): @param doc: capa result doc """ + # inform model that changes are about to occur self.beginResetModel() for rule in rutils.capability_rules(doc): @@ -445,6 +450,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): self.render_capa_doc_match(parent2, match, doc) + # inform model changes have ended self.endResetModel() def capa_doc_feature_to_display(self, feature): @@ -471,7 +477,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): @param parent: parent node to which child is assigned @param feature: capa doc feature node - @param location: locations identified for feature + @param locations: locations identified for feature @param doc: capa doc Example: @@ -484,10 +490,12 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): display = self.capa_doc_feature_to_display(feature) if len(locations) == 1: + # only one location for feature so no need to nest children parent2 = self.render_capa_doc_feature(parent, feature, next(iter(locations)), doc, display=display) else: # feature has multiple children, nest under one parent feature node parent2 = CapaExplorerFeatureItem(parent, display) + for location in sorted(locations): self.render_capa_doc_feature(parent2, feature, location, doc) @@ -509,11 +517,6 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): "type": "number" } """ - instruction_view = ("bytes", "api", "mnemonic", "number", "offset") - byte_view = ("section",) - string_view = ("string",) - default_feature_view = ("import", "export") - # special handling for characteristic pending type if feature["type"] == "characteristic": if feature[feature["type"]] in ("embedded pe",): @@ -522,44 +525,52 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel): if feature[feature["type"]] in ("loop", "recursive call", "tight loop", "switch"): return CapaExplorerFeatureItem(parent, display=display) - # default to instruction view + # default to instruction view for all other characteristics return CapaExplorerInstructionViewItem(parent, display, location) if feature["type"] == "match": + # display content of rule for all rule matches return CapaExplorerRuleMatchItem( parent, display, source=doc["rules"].get(feature[feature["type"]], {}).get("source", "") ) - if feature["type"] in instruction_view: + if feature["type"] in ("bytes", "api", "mnemonic", "number", "offset"): + # display instruction preview return CapaExplorerInstructionViewItem(parent, display, location) - if feature["type"] in byte_view: + if feature["type"] in ("section",): + # display byte preview return CapaExplorerByteViewItem(parent, display, location) - if feature["type"] in string_view: + if feature["type"] in ("string",): + # display string preview return CapaExplorerStringViewItem(parent, display, location) - if feature["type"] in default_feature_view: + if feature["type"] in ("import", "export"): + # display no preview return CapaExplorerFeatureItem(parent, display=display) raise RuntimeError("unexpected feature type: " + str(feature["type"])) def update_function_name(self, old_name, new_name): - """ update all instances of function name + """ update all instances of old function name with new function name @param old_name: previous function name @param new_name: new function name """ + # create empty root index for search root_index = self.index(0, 0, QtCore.QModelIndex()) - # convert name to view format for matching + # convert name to view format for matching e.g. function(my_function) old_name = CapaExplorerFunctionItem.fmt % old_name + # recursive search for all instances of old function name for model_index in self.match( root_index, QtCore.Qt.DisplayRole, old_name, hits=-1, flags=QtCore.Qt.MatchRecursive ): if not isinstance(model_index.internalPointer(), CapaExplorerFunctionItem): continue + # replace old function name with new function name and emit change model_index.internalPointer().info = new_name self.dataChanged.emit(model_index, model_index) diff --git a/capa/ida/explorer/proxy.py b/capa/ida/explorer/proxy.py index 9ebc4eb2..41e962f2 100644 --- a/capa/ida/explorer/proxy.py +++ b/capa/ida/explorer/proxy.py @@ -34,6 +34,7 @@ class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel): def filterAcceptsRow(self, row, parent): """ true if the item in the row indicated by the given row and parent should be included in the model; otherwise returns false + @param row: int @param parent: QModelIndex* diff --git a/capa/ida/explorer/view.py b/capa/ida/explorer/view.py index ce8eba8b..b5273640 100644 --- a/capa/ida/explorer/view.py +++ b/capa/ida/explorer/view.py @@ -17,18 +17,11 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): """ def __init__(self, model, parent=None): - """ initialize CapaExplorerQTreeView - - TODO - - @param model: TODO - @param parent: TODO - """ + """ initialize CapaExplorerQTreeView """ super(CapaExplorerQtreeView, self).__init__(parent) self.setModel(model) - # TODO: get from parent?? self.model = model self.parent = parent @@ -49,7 +42,6 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): # connect slots self.customContextMenuRequested.connect(self.slot_custom_context_menu_requested) self.doubleClicked.connect(self.slot_double_click) - # self.clicked.connect(self.slot_click) self.setStyleSheet("QTreeView::item {padding-right: 15 px;padding-bottom: 2 px;}") @@ -63,10 +55,7 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): self.resize_columns_to_content() def resize_columns_to_content(self): - """ reset view columns to contents - - TODO: prevent columns from shrinking - """ + """ reset view columns to contents """ self.header().resizeSections(QtWidgets.QHeaderView.ResizeToContents) def map_index_to_source_item(self, model_index): @@ -109,10 +98,10 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @yield QAction* """ - default_actions = [ + default_actions = ( ("Copy column", data, self.slot_copy_column), ("Copy row", data, self.slot_copy_row), - ] + ) # add default actions for action in default_actions: @@ -125,9 +114,9 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @yield QAction* """ - function_actions = [ + function_actions = ( ("Rename function", data, self.slot_rename_function), - ] + ) # add function actions for action in function_actions: @@ -180,10 +169,8 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): @param menu: TODO @param pos: TODO """ - if not menu: - return - - menu.exec_(self.viewport().mapToGlobal(pos)) + if menu: + menu.exec_(self.viewport().mapToGlobal(pos)) def slot_copy_column(self, action): """ slot connected to custom context menu @@ -199,7 +186,7 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): def slot_copy_row(self, action): """ slot connected to custom context menu - allows user to select a row and copy the space-delimeted + allows user to select a row and copy the space-delimited data to clipboard @param action: QAction* @@ -249,10 +236,6 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): # show custom context menu at view position self.show_custom_context_menu(menu, pos) - def slot_click(self): - """ slot connected to single click event """ - pass - def slot_double_click(self, model_index): """ slot connected to double click event @@ -264,16 +247,10 @@ class CapaExplorerQtreeView(QtWidgets.QTreeView): item = self.map_index_to_source_item(model_index) column = model_index.column() - if CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS == column: + if CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS == column and item.location: # user double-clicked virtual address column - navigate IDA to address - try: - idc.jumpto(int(item.data(1), 16)) - except ValueError: - pass + idc.jumpto(item.location) if CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column: # user double-clicked information column - un/expand - if self.isExpanded(model_index): - self.collapse(model_index) - else: - self.expand(model_index) + self.collapse(model_index) if self.isExpanded(model_index) else self.expand(model_index) diff --git a/capa/ida/ida_capa_explorer.py b/capa/ida/ida_capa_explorer.py index 970e8a62..980bf945 100644 --- a/capa/ida/ida_capa_explorer.py +++ b/capa/ida/ida_capa_explorer.py @@ -1,15 +1,16 @@ import os import logging import collections +from PyQt5 import QtGui, QtCore, QtWidgets import idaapi -from PyQt5 import QtGui, QtCore, QtWidgets import capa.main import capa.rules import capa.ida.helpers import capa.render.utils as rutils import capa.features.extractors.ida + from capa.ida.explorer.view import CapaExplorerQtreeView from capa.ida.explorer.model import CapaExplorerDataModel from capa.ida.explorer.proxy import CapaExplorerSortFilterProxyModel @@ -23,8 +24,8 @@ class CapaExplorerIdaHooks(idaapi.UI_Hooks): def __init__(self, screen_ea_changed_hook, action_hooks): """ facilitate IDA UI hooks - @param screen_ea_changed: TODO - @param action_hooks: TODO + @param screen_ea_changed_hook: function hook for IDA screen ea changed + @param action_hooks: dict of IDA action handles """ super(CapaExplorerIdaHooks, self).__init__() @@ -76,8 +77,9 @@ class CapaExplorerForm(idaapi.PluginForm): super(CapaExplorerForm, self).__init__() self.form_title = PLUGIN_NAME - self.parent = None self.file_loc = __file__ + + self.parent = None self.ida_hooks = None # models @@ -145,17 +147,17 @@ class CapaExplorerForm(idaapi.PluginForm): self.load_view_parent() def load_view_tabs(self): - """ """ + """ load tabs """ tabs = QtWidgets.QTabWidget() self.view_tabs = tabs def load_view_menu_bar(self): - """ """ + """ load menu bar """ bar = QtWidgets.QMenuBar() self.view_menu_bar = bar def load_view_summary(self): - """ """ + """ load capa summary table """ table_headers = [ "Capability", "Namespace", @@ -177,7 +179,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_summary = table def load_view_attack(self): - """ """ + """ load MITRE ATT&CK table """ table_headers = [ "ATT&CK Tactic", "ATT&CK Technique ", @@ -199,12 +201,12 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_attack = table def load_view_checkbox_limit_by(self): - """ """ + """ load limit results by function checkbox """ check = QtWidgets.QCheckBox("Limit results to current function") check.setChecked(False) check.stateChanged.connect(self.slot_checkbox_limit_by_changed) - self.view_checkbox_limit_by = check + self.view_limit_results_by_function = check def load_view_parent(self): """ load view parent """ @@ -216,9 +218,9 @@ class CapaExplorerForm(idaapi.PluginForm): self.parent.setLayout(layout) def load_view_tree_tab(self): - """ load view tree tab """ + """ load capa tree tab view """ layout = QtWidgets.QVBoxLayout() - layout.addWidget(self.view_checkbox_limit_by) + layout.addWidget(self.view_limit_results_by_function) layout.addWidget(self.view_tree) tab = QtWidgets.QWidget() @@ -227,7 +229,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_tabs.addTab(tab, "Tree View") def load_view_summary_tab(self): - """ """ + """ load capa summary tab view """ layout = QtWidgets.QVBoxLayout() layout.addWidget(self.view_summary) @@ -237,7 +239,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_tabs.addTab(tab, "Summary") def load_view_attack_tab(self): - """ """ + """ load MITRE ATT&CK tab view """ layout = QtWidgets.QVBoxLayout() layout.addWidget(self.view_attack) @@ -255,14 +257,13 @@ class CapaExplorerForm(idaapi.PluginForm): menu = self.view_menu_bar.addMenu("File") - for name, _, handle in actions: + for (name, _, handle) in actions: action = QtWidgets.QAction(name, self.parent) action.triggered.connect(handle) - # action.setToolTip(tip) menu.addAction(action) def load_ida_hooks(self): - """ """ + """ load IDA Pro UI hooks """ action_hooks = { "MakeName": self.ida_hook_rename, "EditFunction": self.ida_hook_rename, @@ -272,7 +273,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.ida_hooks.hook() def unload_ida_hooks(self): - """ unhook IDA user interface """ + """ unload IDA Pro UI hooks """ if self.ida_hooks: self.ida_hooks.unhook() @@ -282,8 +283,8 @@ class CapaExplorerForm(idaapi.PluginForm): called twice, once before action and once after action completes - @param meta: TODO - @param post: TODO + @param meta: metadata cache + @param post: indicates pre or post action """ location = idaapi.get_screen_ea() if not location or not capa.ida.helpers.is_func_start(location): @@ -299,8 +300,13 @@ class CapaExplorerForm(idaapi.PluginForm): meta["prev_name"] = curr_name def ida_hook_screen_ea_changed(self, widget, new_ea, old_ea): - """ """ - if not self.view_checkbox_limit_by.isChecked(): + """ hook for IDA screen ea changed + + @param widget: IDA widget type + @param new_ea: destination ea + @param old_ea: source ea + """ + if not self.view_limit_results_by_function.isChecked(): # ignore if checkbox not selected return @@ -328,7 +334,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_tree.resize_columns_to_content() def load_capa_results(self): - """ """ + """ run capa analysis and render results in UI """ logger.info("-" * 80) logger.info(" Using default embedded rules.") logger.info(" ") @@ -376,11 +382,14 @@ class CapaExplorerForm(idaapi.PluginForm): logger.info("render views completed.") def set_view_tree_default_sort_order(self): - """ """ + """ set capa tree view default sort order """ self.view_tree.sortByColumn(CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION, QtCore.Qt.AscendingOrder) def render_capa_doc_summary(self, doc): - """ """ + """ render capa summary results + + @param doc: capa doc + """ for (row, rule) in enumerate(rutils.capability_rules(doc)): count = len(rule["matches"]) @@ -398,7 +407,10 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_summary.resizeColumnsToContents() def render_capa_doc_mitre_summary(self, doc): - """ """ + """ render capa MITRE ATT&CK results + + @param doc: capa doc + """ tactics = collections.defaultdict(set) for rule in rutils.capability_rules(doc): @@ -418,8 +430,9 @@ class CapaExplorerForm(idaapi.PluginForm): column_one = [] column_two = [] - for tactic, techniques in sorted(tactics.items()): + for (tactic, techniques) in sorted(tactics.items()): column_one.append(tactic.upper()) + # add extra space when more than one technique column_one.extend(["" for i in range(len(techniques) - 1)]) for spec in sorted(techniques): @@ -444,7 +457,7 @@ class CapaExplorerForm(idaapi.PluginForm): self.view_attack.resizeColumnsToContents() def render_new_table_header_item(self, text): - """ """ + """ create new table header item with default style """ item = QtWidgets.QTableWidgetItem(text) item.setForeground(QtGui.QColor(88, 139, 174)) @@ -456,10 +469,10 @@ class CapaExplorerForm(idaapi.PluginForm): return item def ida_reset(self): - """ reset IDA user interface """ + """ reset IDA UI """ self.model_data.reset() self.view_tree.reset() - self.view_checkbox_limit_by.setChecked(False) + self.view_limit_results_by_function.setChecked(False) self.set_view_tree_default_sort_order() def reload(self): @@ -474,7 +487,7 @@ class CapaExplorerForm(idaapi.PluginForm): idaapi.info("%s reload completed." % PLUGIN_NAME) def reset(self): - """ reset user interface elements + """ reset UI elements e.g. checkboxes and IDA highlighting """ @@ -501,10 +514,11 @@ class CapaExplorerForm(idaapi.PluginForm): in function, otherwise clear filter """ match = "" - if self.view_checkbox_limit_by.isChecked(): + if self.view_limit_results_by_function.isChecked(): location = capa.ida.helpers.get_func_start_ea(idaapi.get_screen_ea()) if location: match = capa.ida.explorer.item.location_to_hex(location) + self.model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match) self.view_tree.resize_columns_to_content()