From 18c87e4e55a7b50a0037bea0eda224b0ec3073db Mon Sep 17 00:00:00 2001
From: Moritz Raabe <moritz.raabe@fireeye.com>
Date: Mon, 28 Jun 2021 14:11:28 +0200
Subject: [PATCH] ida extract library funcs identified via flirt

---
 CHANGELOG.md                         |  3 ++-
 capa/features/extractors/ida/insn.py | 21 ++++++++++++++++++---
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 377242ab..4165571f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,7 +18,6 @@ It includes many new rules, including all new techniques introduced in MITRE ATT
 - show-features: don't show features from library functions #569 @williballenthin
 - linter: summarize results at the end #571 @williballenthin
 - linter: check for `or` with always true child statement, e.g. `optional`, colors #348 @mr-tz
-- explorer: add argument to control whether to automatically analyze when running capa explorer #548 @Ana06
 
 ### Breaking Changes
 
@@ -138,6 +137,8 @@ It includes many new rules, including all new techniques introduced in MITRE ATT
 - explorer: document IDA 7.6sp1 as alternative to the patch #536 @Ana06
 - explorer: add support for function-name feature #618 @mike-hunhoff
 - explorer: circular import workaround #654 @mike-hunhoff
+- explorer: add argument to control whether to automatically analyze when running capa explorer #548 @Ana06
+- explorer: extract API features via function names recognized by IDA/FLIRT #661 @mr-tz
 
 ### Development
 
diff --git a/capa/features/extractors/ida/insn.py b/capa/features/extractors/ida/insn.py
index 044e5923..0cc5fc72 100644
--- a/capa/features/extractors/ida/insn.py
+++ b/capa/features/extractors/ida/insn.py
@@ -54,9 +54,6 @@ def get_imports(ctx):
 
 def check_for_api_call(ctx, insn):
     """check instruction for API call"""
-    if not insn.get_canon_mnem() in ("call", "jmp"):
-        return
-
     info = ()
     ref = insn.ea
 
@@ -95,11 +92,29 @@ def extract_insn_api_features(f, bb, insn):
     example:
         call dword [0x00473038]
     """
+    if not insn.get_canon_mnem() in ("call", "jmp"):
+        return
+
     for api in check_for_api_call(f.ctx, insn):
         dll, _, symbol = api.rpartition(".")
         for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
             yield API(name), insn.ea
 
+    # extract IDA/FLIRT recognized API functions
+    targets = list(idautils.CodeRefsFrom(insn.ea, False))
+    if not targets:
+        return
+
+    target = targets[0]
+    target_func = idaapi.get_func(target)
+    if not target_func or target_func.start_ea != target:
+        # not a function (start)
+        return
+
+    if idaapi.get_func(target).flags & idaapi.FUNC_LIB:
+        name = idaapi.get_name(target)
+        yield API(name), insn.ea
+
 
 def extract_insn_number_features(f, bb, insn):
     """parse instruction number features