From 090ec46ca4e6aec6959fc0600f4d6fc687156cec Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Mon, 31 Aug 2020 17:13:10 -0600
Subject: [PATCH] features: extract import A/W variants and their base names

closes #246
---
 capa/features/extractors/helpers.py  | 29 ++++++++++++++++++++++++++++
 capa/features/extractors/ida/file.py | 14 ++++++++++----
 capa/features/extractors/viv/file.py | 11 +++++------
 tests/fixtures.py                    |  2 ++
 4 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/capa/features/extractors/helpers.py b/capa/features/extractors/helpers.py
index 27fba835..2b1c66d6 100644
--- a/capa/features/extractors/helpers.py
+++ b/capa/features/extractors/helpers.py
@@ -9,6 +9,7 @@
 import sys
 import builtins
 
+from capa.features.file import Import
 from capa.features.insn import API
 
 MIN_STACKSTRING_LEN = 8
@@ -65,6 +66,34 @@ def generate_api_features(apiname, va):
             yield API(impname[:-1]), va
 
 
+def is_ordinal(symbol):
+    return symbol[0] == "#"
+
+
+def generate_import_features(dll, symbol, va):
+    """
+    for a given dll, symbol, and address, generate import features.
+    we over-generate features to make matching easier.
+    these include:
+      - kernel32.CreateFileA
+      - kernel32.CreateFile
+      - CreateFileA
+      - CreateFile
+    """
+    # (kernel32.CreateFileA, 0x401000)
+    yield Import(dll + "." + symbol), va
+    # (CreateFileA, 0x401000)
+    if not is_ordinal(symbol):
+        yield Import(symbol), va
+
+    if is_aw_function(symbol):
+        # (kernel32.CreateFile, 0x401000)
+        yield Import(dll + "." + symbol[:-1]), va
+        # (CreateFile, 0x401000)
+        if not is_ordinal(symbol):
+            yield Import(symbol[:-1]), va
+
+
 def all_zeros(bytez):
     return all(b == 0 for b in builtins.bytes(bytez))
 
diff --git a/capa/features/extractors/ida/file.py b/capa/features/extractors/ida/file.py
index 2acc398b..6721344c 100644
--- a/capa/features/extractors/ida/file.py
+++ b/capa/features/extractors/ida/file.py
@@ -97,10 +97,16 @@ def extract_file_import_names():
     """
     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
+            dll = info[0]
+            symbol = info[1]
+        elif info[2]:
+            dll = info[0]
+            symbol = "#%d" % (info[2])
+        else:
+            continue
+
+        for feature, ea in capa.features.extractors.helpers.generate_import_features(dll, symbol, ea):
+            yield feature, ea
 
 
 def extract_file_section_names():
diff --git a/capa/features/extractors/viv/file.py b/capa/features/extractors/viv/file.py
index f0b4d6db..32086e28 100644
--- a/capa/features/extractors/viv/file.py
+++ b/capa/features/extractors/viv/file.py
@@ -9,8 +9,9 @@
 import PE.carve as pe_carve  # vivisect PE
 
 import capa.features.extractors.strings
+import capa.features.extractors.helpers
 from capa.features import String, Characteristic
-from capa.features.file import Export, Import, Section
+from capa.features.file import Export, Section
 
 
 def extract_file_embedded_pe(vw, file_path):
@@ -41,11 +42,9 @@ def extract_file_import_names(vw, file_path):
         if is_viv_ord_impname(impname):
             # replace ord prefix with #
             impname = "#%s" % impname[len("ord") :]
-            tinfo = "%s.%s" % (modname, impname)
-            yield Import(tinfo), va
-        else:
-            yield Import(tinfo), va
-            yield Import(impname), va
+
+        for feature, va in capa.features.extractors.helpers.generate_import_features(modname, impname, va):
+            yield feature, va
 
 
 def is_viv_ord_impname(impname):
diff --git a/tests/fixtures.py b/tests/fixtures.py
index e1f33225..2351730f 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -283,6 +283,8 @@ FEATURE_PRESENCE_TESTS = [
     ("mimikatz", "file", capa.features.file.Import("nope"), False),
     ("mimikatz", "file", capa.features.file.Import("advapi32.CryptAcquireContextW"), True),
     ("mimikatz", "file", capa.features.file.Import("advapi32.CryptAcquireContext"), True),
+    ("mimikatz", "file", capa.features.file.Import("CryptAcquireContextW"), True),
+    ("mimikatz", "file", capa.features.file.Import("CryptAcquireContext"), True),
     # function/characteristic(loop)
     ("mimikatz", "function=0x401517", capa.features.Characteristic("loop"), True),
     ("mimikatz", "function=0x401000", capa.features.Characteristic("loop"), False),