fix symbol generation, ordinals

capa/features/extractors/helpers.py

@@ -50,9 +50,10 @@ def generate_symbols(dll: str, symbol: str, include_dll=False) -> Iterator[str]:
       - CreateFile
       - ws2_32.#1
 
-    note that since capa v7 only `import` features include DLL names:
+    note that since capa v7 only `import` features and APIs called via ordinal include DLL names:
       - kernel32.CreateFileA
       - kernel32.CreateFile
+      - ws2_32.#1
 
     for `api` features dll names are good for documentation but not used during matching
     """
@@ -63,7 +64,7 @@ def generate_symbols(dll: str, symbol: str, include_dll=False) -> Iterator[str]:
     dll = dll[0:-4] if dll.endswith(".dll") else dll
     dll = dll[0:-4] if dll.endswith(".drv") else dll
 
-    if include_dll:
+    if include_dll or is_ordinal(symbol):
         # ws2_32.#1
         # kernel32.CreateFileA
         yield f"{dll}.{symbol}"
@@ -72,11 +73,11 @@ def generate_symbols(dll: str, symbol: str, include_dll=False) -> Iterator[str]:
         # CreateFileA
         yield symbol
 
-        if include_dll:
-            # kernel32.CreateFile
-            yield f"{dll}.{symbol[:-1]}"
-
         if is_aw_function(symbol):
+            if include_dll:
+                # kernel32.CreateFile
+                yield f"{dll}.{symbol[:-1]}"
+
             # CreateFile
             yield symbol[:-1]
 

capa/rules/__init__.py

@@ -596,6 +596,10 @@ def pop_statement_description_entry(d):
 
 
 def trim_dll_part(api: str) -> str:
+    # ordinal imports, like ws2_32.#1, keep dll
+    if ".#" in api:
+        return api
+
     # kernel32.CreateFileA
     if api.count(".") == 1:
         api = api.split(".")[1]

tests/test_helpers.py

@@ -20,3 +20,47 @@ def test_all_zeros():
     assert helpers.all_zeros(b) is True
     assert helpers.all_zeros(c) is False
     assert helpers.all_zeros(d) is False
+
+
+def test_generate_symbols():
+    assert list(helpers.generate_symbols("name.dll", "api", include_dll=True)) == list(
+        helpers.generate_symbols("name", "api", include_dll=True)
+    )
+    assert list(helpers.generate_symbols("name.dll", "api", include_dll=False)) == list(
+        helpers.generate_symbols("name", "api", include_dll=False)
+    )
+
+    # A/W import
+    symbols = list(helpers.generate_symbols("kernel32", "CreateFileA", include_dll=True))
+    assert len(symbols) == 4
+    assert "kernel32.CreateFileA" in symbols
+    assert "kernel32.CreateFile" in symbols
+    assert "CreateFileA" in symbols
+    assert "CreateFile" in symbols
+
+    # import
+    symbols = list(helpers.generate_symbols("kernel32", "WriteFile", include_dll=True))
+    assert len(symbols) == 2
+    assert "kernel32.WriteFile" in symbols
+    assert "WriteFile" in symbols
+
+    # ordinal import
+    symbols = list(helpers.generate_symbols("ws2_32", "#1", include_dll=True))
+    assert len(symbols) == 1
+    assert "ws2_32.#1" in symbols
+
+    # A/W api
+    symbols = list(helpers.generate_symbols("kernel32", "CreateFileA", include_dll=False))
+    assert len(symbols) == 2
+    assert "CreateFileA" in symbols
+    assert "CreateFile" in symbols
+
+    # api
+    symbols = list(helpers.generate_symbols("kernel32", "WriteFile", include_dll=False))
+    assert len(symbols) == 1
+    assert "WriteFile" in symbols
+
+    # ordinal api
+    symbols = list(helpers.generate_symbols("ws2_32", "#1", include_dll=False))
+    assert len(symbols) == 1
+    assert "ws2_32.#1" in symbols