diff --git a/capa/features/extractors/helpers.py b/capa/features/extractors/helpers.py index 71d28ef52..c4d4acb6f 100644 --- a/capa/features/extractors/helpers.py +++ b/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] diff --git a/capa/rules/__init__.py b/capa/rules/__init__.py index be04ec557..8cfbab968 100644 --- a/capa/rules/__init__.py +++ b/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] diff --git a/tests/test_helpers.py b/tests/test_helpers.py index f956be5cb..56a511b67 100644 --- a/tests/test_helpers.py +++ b/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