From 3772c5c0bccf48f00f76dd994655b20622fb4e7e Mon Sep 17 00:00:00 2001
From: Michael Hunhoff <mike.hunhoff@gmail.com>
Date: Thu, 27 Aug 2020 12:07:44 -0600
Subject: [PATCH] add additional nzxor stack cookie check for IDA extractor

---
 capa/features/extractors/ida/__init__.py |  2 +-
 capa/features/extractors/ida/helpers.py  | 18 +++++++++++++++
 capa/features/extractors/ida/insn.py     | 29 ++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/capa/features/extractors/ida/__init__.py b/capa/features/extractors/ida/__init__.py
index 6faf9847..362b0da7 100644
--- a/capa/features/extractors/ida/__init__.py
+++ b/capa/features/extractors/ida/__init__.py
@@ -75,7 +75,7 @@ class IdaFeatureExtractor(FeatureExtractor):
             yield feature, ea
 
     def get_basic_blocks(self, f):
-        for bb in idaapi.FlowChart(f, flags=idaapi.FC_PREDS):
+        for bb in capa.features.extractors.ida.helpers.get_function_blocks(f):
             yield add_ea_int_cast(bb)
 
     def extract_basic_block_features(self, f, bb):
diff --git a/capa/features/extractors/ida/helpers.py b/capa/features/extractors/ida/helpers.py
index cba352a9..22a99f71 100644
--- a/capa/features/extractors/ida/helpers.py
+++ b/capa/features/extractors/ida/helpers.py
@@ -341,3 +341,21 @@ def find_data_reference_from_insn(insn, max_depth=10):
         ea = data_refs[0]
 
     return ea
+
+
+def get_function_blocks(f):
+    """yield basic blocks contained in specified function
+
+    args:
+        f (IDA func_t)
+    yield:
+        block (IDA BasicBlock)
+    """
+    # leverage idaapi.FC_NOEXT flag to ignore useless external blocks referenced by the function
+    for block in idaapi.FlowChart(f, flags=(idaapi.FC_PREDS | idaapi.FC_NOEXT)):
+        yield block
+
+
+def is_basic_block_return(bb):
+    """ check if basic block is return block """
+    return bb.type == idaapi.fcb_ret
diff --git a/capa/features/extractors/ida/insn.py b/capa/features/extractors/ida/insn.py
index 80844788..73f98056 100644
--- a/capa/features/extractors/ida/insn.py
+++ b/capa/features/extractors/ida/insn.py
@@ -15,6 +15,10 @@ import capa.features.extractors.ida.helpers
 from capa.features import ARCH_X32, ARCH_X64, MAX_BYTES_FEATURE_SIZE, Bytes, String, Characteristic
 from capa.features.insn import Number, Offset, Mnemonic
 
+# security cookie checks may perform non-zeroing XORs, these are expected within a certain
+# byte range within the first and returning basic blocks, this helps to reduce FP features
+SECURITY_COOKIE_BYTES_DELTA = 0x40
+
 
 def get_arch(ctx):
     """
@@ -223,12 +227,37 @@ def bb_stack_cookie_registers(bb):
                     yield op.reg
 
 
+def is_nzxor_stack_cookie_delta(f, bb, insn):
+    """ check if nzxor exists within stack cookie delta """
+    # security cookie check should use SP or BP
+    if not capa.features.extractors.ida.helpers.is_frame_register(insn.Op2.reg):
+        return False
+
+    f_bbs = tuple(capa.features.extractors.ida.helpers.get_function_blocks(f))
+
+    # expect security cookie init in first basic block within first bytes (instructions)
+    if capa.features.extractors.ida.helpers.is_basic_block_equal(bb, f_bbs[0]) and insn.ea < (
+        bb.start_ea + SECURITY_COOKIE_BYTES_DELTA
+    ):
+        return True
+
+    # ... or within last bytes (instructions) before a return
+    if capa.features.extractors.ida.helpers.is_basic_block_return(bb) and insn.ea > (
+        bb.start_ea + capa.features.extractors.ida.helpers.basic_block_size(bb) - SECURITY_COOKIE_BYTES_DELTA
+    ):
+        return True
+
+    return False
+
+
 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)):
         # Example:
         #   xor     ecx, ebp        ; StackCookie
         return True
+    if is_nzxor_stack_cookie_delta(f, bb, insn):
+        return True
     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: