tests: add feature tests for idalib

tests/fixtures.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import io
+import logging
 import contextlib
 import collections
 from pathlib import Path
@@ -21,6 +22,7 @@ from functools import lru_cache
 import pytest
 
 import capa.main
+import capa.helpers
 import capa.features.file
 import capa.features.insn
 import capa.features.common
@@ -53,6 +55,7 @@ from capa.features.extractors.base_extractor import (
 )
 from capa.features.extractors.dnfile.extractor import DnfileFeatureExtractor
 
+logger = logging.getLogger(__name__)
 CD = Path(__file__).resolve().parent
 DOTNET_DIR = CD / "data" / "dotnet"
 DNFILE_TESTFILES = DOTNET_DIR / "dnfile-testfiles"
@@ -199,6 +202,41 @@ def get_binja_extractor(path: Path):
 
     return extractor
 
+# we can't easily cache this because the extractor relies on global state (the opened database)
+# which also has to be closed elsewhere. so, the idalib tests will just take a little bit to run.
+def get_idalib_extractor(path: Path):
+    import capa.features.extractors.ida.idalib as idalib
+
+    if not idalib.has_idalib():
+        raise RuntimeError("cannot find IDA idalib module.")
+
+    if not idalib.load_idalib():
+        raise RuntimeError("failed to load IDA idalib module.")
+
+    import idapro
+    import ida_auto
+
+    import capa.features.extractors.ida.extractor
+
+    logger.debug("idalib: opening database...")
+
+    idapro.enable_console_messages(False)
+    # - 0 - Success (database not packed)
+    # - 1 - Success (database was packed)
+    # - 2 - User cancelled or 32-64 bit conversion failed
+    # - 4 - Database initialization failed
+    # - -1 - Generic errors (database already open, auto-analysis failed, etc.)
+    # - -2 - User cancelled operation
+    ret = idapro.open_database(str(path), run_auto_analysis=True)
+    if ret not in (0, 1):
+        raise RuntimeError("failed to analyze input file")
+
+    logger.debug("idalib: waiting for analysis...")
+    ida_auto.auto_wait()
+    logger.debug("idalib: opened database.")
+
+    return capa.features.extractors.ida.extractor.IdaFeatureExtractor()
+
 
 @lru_cache(maxsize=1)
 def get_cape_extractor(path):

tests/test_idalib_features.py

@@ -0,0 +1,58 @@
+
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import pytest
+
+import capa.features.extractors.ida.idalib
+
+import fixtures
+
+logger = logging.getLogger(__name__)
+
+idalib_present = capa.features.extractors.ida.idalib.has_idalib()
+
+
+@pytest.mark.skipif(idalib_present is False, reason="Skip idalib tests if the idalib Python API is not installed")
+@fixtures.parametrize(
+    "sample,scope,feature,expected",
+    fixtures.FEATURE_PRESENCE_TESTS + fixtures.FEATURE_SYMTAB_FUNC_TESTS,
+    indirect=["sample", "scope"],
+)
+def test_idalib_features(sample, scope, feature, expected):
+    try:
+        fixtures.do_test_feature_presence(fixtures.get_idalib_extractor, sample, scope, feature, expected)
+    finally:
+        logger.debug("closing database...")
+        import idapro
+        idapro.close_database(save=False)
+        logger.debug("opened database.")
+
+
+@pytest.mark.skipif(idalib_present is False, reason="Skip idalib tests if the idalib Python API is not installed")
+@fixtures.parametrize(
+    "sample,scope,feature,expected",
+    fixtures.FEATURE_COUNT_TESTS,
+    indirect=["sample", "scope"],
+)
+def test_idalib_feature_counts(sample, scope, feature, expected):
+    try:
+        fixtures.do_test_feature_count(fixtures.get_idalib_extractor, sample, scope, feature, expected)
+    finally:
+        logger.debug("closing database...")
+        import idapro
+        idapro.close_database(save=False)
+        logger.debug("closed database.")