From 8cbe3f8546f48870d752b54669b521b23eba619a Mon Sep 17 00:00:00 2001 From: William Ballenthin Date: Fri, 14 Aug 2020 11:25:00 -0600 Subject: [PATCH] tests: move expected features into fixtures for reuse closes #225 --- tests/fixtures.py | 303 +++++++++++++++++++++++++++++++- tests/test_viv_features.py | 341 ++----------------------------------- 2 files changed, 317 insertions(+), 327 deletions(-) diff --git a/tests/fixtures.py b/tests/fixtures.py index a4a59f82..8cd2ef48 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -11,81 +11,380 @@ import os.path import collections import pytest -import viv_utils + +import capa.features.file +import capa.features.insn +import capa.features.basicblock +from capa.features import ARCH_X32, ARCH_X64 + +try: + from functools import lru_cache +except ImportError: + from backports.functools_lru_cache import lru_cache + CD = os.path.dirname(__file__) +@lru_cache() +def extract_file_features(extractor): + features = collections.defaultdict(set) + for feature, va in extractor.extract_file_features(): + features[feature].add(va) + return features + + +@lru_cache() +def extract_function_features(extractor, f): + features = collections.defaultdict(set) + for bb in extractor.get_basic_blocks(f): + for insn in extractor.get_instructions(f, bb): + for feature, va in extractor.extract_insn_features(f, bb, insn): + features[feature].add(va) + for feature, va in extractor.extract_basic_block_features(f, bb): + features[feature].add(va) + for feature, va in extractor.extract_function_features(f): + features[feature].add(va) + return features + + +@lru_cache() +def extract_basic_block_features(extractor, f, bb): + features = collections.defaultdict(set) + for insn in extractor.get_instructions(f, bb): + for feature, va in extractor.extract_insn_features(f, bb, insn): + features[feature].add(va) + for feature, va in extractor.extract_basic_block_features(f, bb): + features[feature].add(va) + return features + + +@pytest.fixture +def sample(request): + if request.param == "mimikatz": + return os.path.join(CD, "data", "mimikatz.exe_") + elif request.param == "kernel32": + return os.path.join(CD, "data", "kernel32.dll_") + elif request.param == "kernel32-64": + return os.path.join(CD, "data", "kernel32-64.dll_") + elif request.param == "pma12-04": + return os.path.join(CD, "data", "Practical Malware Analysis Lab 12-04.exe_") + elif request.param.startswith("a1982"): + return os.path.join(CD, "data", "a198216798ca38f280dc413f8c57f2c2.exe_") + elif request.param.startswith("39c05"): + return os.path.join(CD, "data", "39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.dll_") + elif request.param.startswith("c9188"): + return os.path.join(CD, "data", "c91887d861d9bd4a5872249b641bc9f9.exe_") + else: + raise ValueError("unexpected sample fixture") + + +def get_function(extractor, fva): + for f in extractor.get_functions(): + if f.__int__() == fva: + return f + raise ValueError("function not found") + + +def get_basic_block(extractor, f, va): + for bb in extractor.get_basic_blocks(f): + if bb.__int__() == va: + return bb + raise ValueError("basic block not found") + + +@pytest.fixture +def scope(request): + if request.param == "file": + + def inner(extractor): + return extract_file_features(extractor) + + inner.__name__ = request.param + return inner + elif "bb=" in request.param: + # like `function=0x401000,bb=0x40100A` + fspec, _, bbspec = request.param.partition(",") + fva = int(fspec.partition("=")[2], 0x10) + bbva = int(bbspec.partition("=")[2], 0x10) + + def inner(extractor): + f = get_function(extractor, fva) + bb = get_basic_block(extractor, f, bbva) + return extract_basic_block_features(extractor, f, bb) + + inner.__name__ = request.param + return inner + elif request.param.startswith("function"): + # like `function=0x401000` + va = int(request.param.partition("=")[2], 0x10) + + def inner(extractor): + f = get_function(extractor, va) + return extract_function_features(extractor, f) + + inner.__name__ = request.param + return inner + else: + raise ValueError("unexpected scope fixture") + + +def parametrize(params, values, **kwargs): + """ + extend `pytest.mark.parametrize` to pretty-print features. + by default, it renders objects as an opaque value. + ref: https://docs.pytest.org/en/2.9.0/example/parametrize.html#different-options-for-test-ids + rendered ID might look something like: + mimikatz-function=0x403BAC-api(CryptDestroyKey)-True + """ + ids = ["-".join(map(str, vs)) for vs in values] + return pytest.mark.parametrize(params, values, ids=ids, **kwargs) + + +FEATURE_PRESENCE_TESTS = [ + # file/characteristic("embedded pe") + ("pma12-04", "file", capa.features.Characteristic("embedded pe"), True), + # file/string + ("mimikatz", "file", capa.features.String("SCardControl"), True), + ("mimikatz", "file", capa.features.String("SCardTransmit"), True), + ("mimikatz", "file", capa.features.String("ACR > "), True), + ("mimikatz", "file", capa.features.String("nope"), False), + # file/sections + ("mimikatz", "file", capa.features.file.Section(".rsrc"), True), + ("mimikatz", "file", capa.features.file.Section(".text"), True), + ("mimikatz", "file", capa.features.file.Section(".nope"), False), + # file/exports + ("kernel32", "file", capa.features.file.Export("BaseThreadInitThunk"), True), + ("kernel32", "file", capa.features.file.Export("lstrlenW"), True), + ("kernel32", "file", capa.features.file.Export("nope"), False), + # file/imports + ("mimikatz", "file", capa.features.file.Import("advapi32.CryptSetHashParam"), True), + ("mimikatz", "file", capa.features.file.Import("CryptSetHashParam"), True), + ("mimikatz", "file", capa.features.file.Import("kernel32.IsWow64Process"), True), + ("mimikatz", "file", capa.features.file.Import("msvcrt.exit"), True), + ("mimikatz", "file", capa.features.file.Import("cabinet.#11"), True), + ("mimikatz", "file", capa.features.file.Import("#11"), False), + ("mimikatz", "file", capa.features.file.Import("#nope"), False), + ("mimikatz", "file", capa.features.file.Import("nope"), False), + # function/characteristic(loop) + ("mimikatz", "function=0x401517", capa.features.Characteristic("loop"), True), + ("mimikatz", "function=0x401000", capa.features.Characteristic("loop"), False), + # bb/characteristic(tight loop) + ("mimikatz", "function=0x402EC4", capa.features.Characteristic("tight loop"), True), + ("mimikatz", "function=0x401000", capa.features.Characteristic("tight loop"), False), + # bb/characteristic(stack string) + ("mimikatz", "function=0x4556E5", capa.features.Characteristic("stack string"), True), + ("mimikatz", "function=0x401000", capa.features.Characteristic("stack string"), False), + # bb/characteristic(tight loop) + ("mimikatz", "function=0x402EC4,bb=0x402F8E", capa.features.Characteristic("tight loop"), True), + ("mimikatz", "function=0x401000,bb=0x401000", capa.features.Characteristic("tight loop"), False), + # insn/mnemonic + ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("push"), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("movzx"), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("xor"), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("in"), False), + ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("out"), False), + # insn/number + ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Number(0x3136B0), True), + # insn/number: stack adjustments + ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xC), False), + ("mimikatz", "function=0x40105D", capa.features.insn.Number(0x10), False), + # insn/number: arch flavors + ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF, arch=ARCH_X32), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF, arch=ARCH_X64), False), + # insn/offset + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x4), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0xC), True), + # insn/offset: stack references + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x8), False), + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x10), False), + # insn/offset: negative + ("mimikatz", "function=0x4011FB", capa.features.insn.Offset(-0x1), True), + ("mimikatz", "function=0x4011FB", capa.features.insn.Offset(-0x2), True), + # insn/offset: arch flavors + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0, arch=ARCH_X32), True), + ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0, arch=ARCH_X64), False), + # insn/api + ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptAcquireContextW"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptAcquireContext"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptGenKey"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptImportKey"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptDestroyKey"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptAcquireContextW"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptAcquireContext"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptGenKey"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptImportKey"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptDestroyKey"), True), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("Nope"), False), + ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.Nope"), False), + # insn/api: thunk + ("mimikatz", "function=0x4556E5", capa.features.insn.API("advapi32.LsaQueryInformationPolicy"), True), + ("mimikatz", "function=0x4556E5", capa.features.insn.API("LsaQueryInformationPolicy"), True), + # insn/api: x64 + ("kernel32-64", "function=0x180001010", capa.features.insn.API("RtlVirtualUnwind"), True,), + ("kernel32-64", "function=0x180001010", capa.features.insn.API("RtlVirtualUnwind"), True), + # insn/api: x64 thunk + ("kernel32-64", "function=0x1800202B0", capa.features.insn.API("RtlCaptureContext"), True,), + ("kernel32-64", "function=0x1800202B0", capa.features.insn.API("RtlCaptureContext"), True), + # insn/api: resolve indirect calls + ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.CreatePipe"), True), + ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.SetHandleInformation"), True), + ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.CloseHandle"), True), + ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.WriteFile"), True), + # insn/string + ("mimikatz", "function=0x40105D", capa.features.String("SCardControl"), True), + ("mimikatz", "function=0x40105D", capa.features.String("SCardTransmit"), True), + ("mimikatz", "function=0x40105D", capa.features.String("ACR > "), True), + ("mimikatz", "function=0x40105D", capa.features.String("nope"), False), + # insn/string, pointer to string + ("mimikatz", "function=0x44EDEF", capa.features.String("INPUTEVENT"), True), + # insn/bytes + ("mimikatz", "function=0x40105D", capa.features.Bytes("SCardControl".encode("utf-16le")), True), + ("mimikatz", "function=0x40105D", capa.features.Bytes("SCardTransmit".encode("utf-16le")), True), + ("mimikatz", "function=0x40105D", capa.features.Bytes("ACR > ".encode("utf-16le")), True), + ("mimikatz", "function=0x40105D", capa.features.Bytes("nope".encode("ascii")), False), + # insn/bytes, pointer to bytes + ("mimikatz", "function=0x44EDEF", capa.features.Bytes("INPUTEVENT".encode("utf-16le")), True), + # insn/characteristic(nzxor) + ("mimikatz", "function=0x410DFC", capa.features.Characteristic("nzxor"), True), + ("mimikatz", "function=0x40105D", capa.features.Characteristic("nzxor"), False), + # insn/characteristic(nzxor): no security cookies + ("mimikatz", "function=0x46D534", capa.features.Characteristic("nzxor"), False), + # insn/characteristic(peb access) + ("kernel32-64", "function=0x1800017D0", capa.features.Characteristic("peb access"), True), + ("mimikatz", "function=0x4556E5", capa.features.Characteristic("peb access"), False), + # insn/characteristic(gs access) + ("kernel32-64", "function=0x180001068", capa.features.Characteristic("gs access"), True), + ("mimikatz", "function=0x4556E5", capa.features.Characteristic("gs access"), False), + # insn/characteristic(cross section flow) + ("a1982...", "function=0x4014D0", capa.features.Characteristic("cross section flow"), True), + # insn/characteristic(cross section flow): imports don't count + ("kernel32-64", "function=0x180001068", capa.features.Characteristic("cross section flow"), False), + ("mimikatz", "function=0x4556E5", capa.features.Characteristic("cross section flow"), False), + # insn/characteristic(recursive call) + ("39c05...", "function=0x10003100", capa.features.Characteristic("recursive call"), True), + ("mimikatz", "function=0x4556E5", capa.features.Characteristic("recursive call"), False), + # insn/characteristic(indirect call) + ("mimikatz", "function=0x4175FF", capa.features.Characteristic("indirect call"), True), + ("mimikatz", "function=0x4556E5", capa.features.Characteristic("indirect call"), False), + # insn/characteristic(calls from) + ("mimikatz", "function=0x4556E5", capa.features.Characteristic("calls from"), True), + ("mimikatz", "function=0x4702FD", capa.features.Characteristic("calls from"), False), + # function/characteristic(calls to) + ("mimikatz", "function=0x40105D", capa.features.Characteristic("calls to"), True), + ("mimikatz", "function=0x46C0D2", capa.features.Characteristic("calls to"), False), +] + +FEATURE_COUNT_TESTS = [ + ("mimikatz", "function=0x40E51B", capa.features.basicblock.BasicBlock(), 1), + ("mimikatz", "function=0x40E5C2", capa.features.basicblock.BasicBlock(), 7), + ("mimikatz", "function=0x40E5C2", capa.features.Characteristic("calls from"), 3), +] + + +def do_test_feature_presence(get_extractor, sample, scope, feature, expected): + extractor = get_extractor(sample) + features = scope(extractor) + if expected: + msg = "%s should be found in %s" % (str(feature), scope.__name__) + else: + msg = "%s should not be found in %s" % (str(feature), scope.__name__) + assert feature.evaluate(features) == expected, msg + + +def do_test_feature_count(get_extractor, sample, scope, feature, expected): + extractor = get_extractor(sample) + features = scope(extractor) + msg = "%s should be found %d times in %s" % (str(feature), expected, scope.__name__) + assert len(features[feature]) == expected, msg + + Sample = collections.namedtuple("Sample", ["vw", "path"]) @pytest.fixture def mimikatz(): + import viv_utils path = os.path.join(CD, "data", "mimikatz.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_a933a1a402775cfa94b6bee0963f4b46(): + import viv_utils path = os.path.join(CD, "data", "a933a1a402775cfa94b6bee0963f4b46.dll_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def kernel32(): + import viv_utils path = os.path.join(CD, "data", "kernel32.dll_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_a198216798ca38f280dc413f8c57f2c2(): + import viv_utils path = os.path.join(CD, "data", "a198216798ca38f280dc413f8c57f2c2.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_9324d1a8ae37a36ae560c37448c9705a(): + import viv_utils path = os.path.join(CD, "data", "9324d1a8ae37a36ae560c37448c9705a.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def pma_lab_12_04(): + import viv_utils path = os.path.join(CD, "data", "Practical Malware Analysis Lab 12-04.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_bfb9b5391a13d0afd787e87ab90f14f5(): + import viv_utils path = os.path.join(CD, "data", "bfb9b5391a13d0afd787e87ab90f14f5.dll_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_lab21_01(): + import viv_utils path = os.path.join(CD, "data", "Practical Malware Analysis Lab 21-01.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_c91887d861d9bd4a5872249b641bc9f9(): + import viv_utils path = os.path.join(CD, "data", "c91887d861d9bd4a5872249b641bc9f9.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41(): + import viv_utils path = os.path.join(CD, "data", "39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.dll_",) return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32(): + import viv_utils path = os.path.join(CD, "data", "499c2a85f6e8142c3f48d4251c9c7cd6.raw32") return Sample(viv_utils.getShellcodeWorkspace(path), path) @pytest.fixture def sample_al_khaser_x86(): + import viv_utils path = os.path.join(CD, "data", "al-khaser_x86.exe_") - return Sample(viv_utils.getWorkspace(path), path) + return Sample(viv_utils.getWorkspace(path), path) \ No newline at end of file diff --git a/tests/test_viv_features.py b/tests/test_viv_features.py index 8eab1249..bc5e2038 100644 --- a/tests/test_viv_features.py +++ b/tests/test_viv_features.py @@ -6,336 +6,27 @@ # 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 viv_utils from fixtures import * -import capa.features -import capa.features.file -import capa.features.insn -import capa.features.basicblock -import capa.features.extractors.viv.file -import capa.features.extractors.viv.insn -import capa.features.extractors.viv.function -import capa.features.extractors.viv.basicblock -from capa.features import ARCH_X32, ARCH_X64 +import capa.main +import capa.features.extractors.viv -def extract_file_features(vw, path): - features = set([]) - for feature, va in capa.features.extractors.viv.file.extract_features(vw, path): - features.add(feature) - return features +@lru_cache() +def get_viv_extractor(path): + vw = capa.main.get_workspace(path, "auto", should_save=True) + return capa.features.extractors.viv.VivisectFeatureExtractor(vw, path) -def extract_function_features(f): - features = collections.defaultdict(set) - for bb in f.basic_blocks: - for insn in bb.instructions: - for feature, va in capa.features.extractors.viv.insn.extract_features(f, bb, insn): - features[feature].add(va) - for feature, va in capa.features.extractors.viv.basicblock.extract_features(f, bb): - features[feature].add(va) - for feature, va in capa.features.extractors.viv.function.extract_features(f): - features[feature].add(va) - return features +@parametrize( + "sample,scope,feature,expected", FEATURE_PRESENCE_TESTS, indirect=["sample", "scope"], +) +def test_viv_features(sample, scope, feature, expected): + do_test_feature_presence(get_viv_extractor, sample, scope, feature, expected) -def extract_basic_block_features(f, bb): - features = set({}) - for insn in bb.instructions: - for feature, _ in capa.features.extractors.viv.insn.extract_features(f, bb, insn): - features.add(feature) - for feature, _ in capa.features.extractors.viv.basicblock.extract_features(f, bb): - features.add(feature) - return features - - -def test_api_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x403BAC)) - assert capa.features.insn.API("advapi32.CryptAcquireContextW") in features - assert capa.features.insn.API("advapi32.CryptAcquireContext") in features - assert capa.features.insn.API("advapi32.CryptGenKey") in features - assert capa.features.insn.API("advapi32.CryptImportKey") in features - assert capa.features.insn.API("advapi32.CryptDestroyKey") in features - assert capa.features.insn.API("CryptAcquireContextW") in features - assert capa.features.insn.API("CryptAcquireContext") in features - assert capa.features.insn.API("CryptGenKey") in features - assert capa.features.insn.API("CryptImportKey") in features - assert capa.features.insn.API("CryptDestroyKey") in features - - -def test_api_features_64_bit(sample_a198216798ca38f280dc413f8c57f2c2): - features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x4011B0)) - assert capa.features.insn.API("kernel32.GetStringTypeA") in features - assert capa.features.insn.API("kernel32.GetStringTypeW") not in features - assert capa.features.insn.API("kernel32.GetStringType") in features - assert capa.features.insn.API("GetStringTypeA") in features - assert capa.features.insn.API("GetStringType") in features - # call via thunk in IDA Pro - features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x401CB0)) - assert capa.features.insn.API("msvcrt.vfprintf") in features - assert capa.features.insn.API("vfprintf") in features - - -def test_string_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D)) - assert capa.features.String("SCardControl") in features - assert capa.features.String("SCardTransmit") in features - assert capa.features.String("ACR > ") in features - # other strings not in this function - assert capa.features.String("bcrypt.dll") not in features - - -def test_string_pointer_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x44EDEF)) - assert capa.features.String("INPUTEVENT") in features - - -def test_byte_features(sample_9324d1a8ae37a36ae560c37448c9705a): - features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60)) - wanted = capa.features.Bytes(b"\xED\x24\x9E\xF4\x52\xA9\x07\x47\x55\x8E\xE1\xAB\x30\x8E\x23\x61") - # use `==` rather than `is` because the result is not `True` but a truthy value. - assert wanted.evaluate(features) == True - - -def test_byte_features64(sample_lab21_01): - features = extract_function_features(viv_utils.Function(sample_lab21_01.vw, 0x1400010C0)) - wanted = capa.features.Bytes(b"\x32\xA2\xDF\x2D\x99\x2B\x00\x00") - # use `==` rather than `is` because the result is not `True` but a truthy value. - assert wanted.evaluate(features) == True - - -def test_bytes_pointer_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x44EDEF)) - assert capa.features.Bytes("INPUTEVENT".encode("utf-16le")).evaluate(features) == True - - -def test_number_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D)) - assert capa.features.insn.Number(0xFF) in features - assert capa.features.insn.Number(0x3136B0) in features - # the following are stack adjustments - assert capa.features.insn.Number(0xC) not in features - assert capa.features.insn.Number(0x10) not in features - - -def test_number_arch_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D)) - assert capa.features.insn.Number(0xFF) in features - assert capa.features.insn.Number(0xFF, arch=ARCH_X32) in features - assert capa.features.insn.Number(0xFF, arch=ARCH_X64) not in features - - -def test_unmapped_immediate_memory_reference_features(sample_al_khaser_x86): - features = extract_function_features(viv_utils.Function(sample_al_khaser_x86.vw, 0x41AAB4)) - assert capa.features.insn.Number(0x7FFE02D4) in features - - -def test_offset_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D)) - assert capa.features.insn.Offset(0x0) in features - assert capa.features.insn.Offset(0x4) in features - assert capa.features.insn.Offset(0xC) in features - # the following are stack references - assert capa.features.insn.Offset(0x8) not in features - assert capa.features.insn.Offset(0x10) not in features - - # this function has the following negative offsets - # movzx ecx, byte ptr [eax-1] - # movzx eax, byte ptr [eax-2] - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x4011FB)) - assert capa.features.insn.Offset(-0x1) in features - assert capa.features.insn.Offset(-0x2) in features - - -def test_offset_arch_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D)) - assert capa.features.insn.Offset(0x0) in features - assert capa.features.insn.Offset(0x0, arch=ARCH_X32) in features - assert capa.features.insn.Offset(0x0, arch=ARCH_X64) not in features - - -def test_nzxor_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x410DFC)) - assert capa.features.Characteristic("nzxor") in features # 0x0410F0B - - -def get_bb_insn(f, va): - """fetch the BasicBlock and Instruction instances for the given VA in the given function.""" - for bb in f.basic_blocks: - for insn in bb.instructions: - if insn.va == va: - return (bb, insn) - raise KeyError(va) - - -def test_is_security_cookie(mimikatz): - # not a security cookie check - f = viv_utils.Function(mimikatz.vw, 0x410DFC) - for va in [0x0410F0B]: - bb, insn = get_bb_insn(f, va) - assert capa.features.extractors.viv.insn.is_security_cookie(f, bb, insn) == False - - # security cookie initial set and final check - f = viv_utils.Function(mimikatz.vw, 0x46C54A) - for va in [0x46C557, 0x46C63A]: - bb, insn = get_bb_insn(f, va) - assert capa.features.extractors.viv.insn.is_security_cookie(f, bb, insn) == True - - -def test_mnemonic_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D)) - assert capa.features.insn.Mnemonic("push") in features - assert capa.features.insn.Mnemonic("movzx") in features - assert capa.features.insn.Mnemonic("xor") in features - - assert capa.features.insn.Mnemonic("in") not in features - assert capa.features.insn.Mnemonic("out") not in features - - -def test_peb_access_features(sample_a933a1a402775cfa94b6bee0963f4b46): - features = extract_function_features(viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA6FEC)) - assert capa.features.Characteristic("peb access") in features - - -def test_file_section_name_features(mimikatz): - features = extract_file_features(mimikatz.vw, mimikatz.path) - assert capa.features.file.Section(".rsrc") in features - assert capa.features.file.Section(".text") in features - assert capa.features.file.Section(".nope") not in features - - -def test_tight_loop_features(mimikatz): - f = viv_utils.Function(mimikatz.vw, 0x402EC4) - for bb in f.basic_blocks: - if bb.va != 0x402F8E: - continue - features = extract_basic_block_features(f, bb) - assert capa.features.Characteristic("tight loop") in features - assert capa.features.basicblock.BasicBlock() in features - - -def test_tight_loop_bb_features(mimikatz): - f = viv_utils.Function(mimikatz.vw, 0x402EC4) - for bb in f.basic_blocks: - if bb.va != 0x402F8E: - continue - features = extract_basic_block_features(f, bb) - assert capa.features.Characteristic("tight loop") in features - assert capa.features.basicblock.BasicBlock() in features - - -def test_file_export_name_features(kernel32): - features = extract_file_features(kernel32.vw, kernel32.path) - assert capa.features.file.Export("BaseThreadInitThunk") in features - assert capa.features.file.Export("lstrlenW") in features - - -def test_file_import_name_features(mimikatz): - features = extract_file_features(mimikatz.vw, mimikatz.path) - assert capa.features.file.Import("advapi32.CryptSetHashParam") in features - assert capa.features.file.Import("CryptSetHashParam") in features - assert capa.features.file.Import("kernel32.IsWow64Process") in features - assert capa.features.file.Import("msvcrt.exit") in features - assert capa.features.file.Import("cabinet.#11") in features - assert capa.features.file.Import("#11") not in features - - -def test_cross_section_flow_features(sample_a198216798ca38f280dc413f8c57f2c2): - features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x4014D0)) - assert capa.features.Characteristic("cross section flow") in features - - # this function has calls to some imports, - # which should not trigger cross-section flow characteristic - features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x401563)) - assert capa.features.Characteristic("cross section flow") not in features - - -def test_segment_access_features(sample_a933a1a402775cfa94b6bee0963f4b46): - features = extract_function_features(viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA6FEC)) - assert capa.features.Characteristic("fs access") in features - - -def test_thunk_features(sample_9324d1a8ae37a36ae560c37448c9705a): - features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x407970)) - assert capa.features.insn.API("kernel32.CreateToolhelp32Snapshot") in features - assert capa.features.insn.API("CreateToolhelp32Snapshot") in features - - -def test_file_embedded_pe(pma_lab_12_04): - features = extract_file_features(pma_lab_12_04.vw, pma_lab_12_04.path) - assert capa.features.Characteristic("embedded pe") in features - - -def test_stackstring_features(mimikatz): - features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x4556E5)) - assert capa.features.Characteristic("stack string") in features - - -def test_recursive_call_feature(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41): - features = extract_function_features( - viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10003100) - ) - assert capa.features.Characteristic("recursive call") in features - - features = extract_function_features( - viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10007B00) - ) - assert capa.features.Characteristic("recursive call") not in features - - -def test_loop_feature(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41): - features = extract_function_features( - viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10003D30) - ) - assert capa.features.Characteristic("loop") in features - - features = extract_function_features( - viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10007250) - ) - assert capa.features.Characteristic("loop") not in features - - -def test_file_string_features(sample_bfb9b5391a13d0afd787e87ab90f14f5): - features = extract_file_features( - sample_bfb9b5391a13d0afd787e87ab90f14f5.vw, sample_bfb9b5391a13d0afd787e87ab90f14f5.path, - ) - assert capa.features.String("WarStop") in features # ASCII, offset 0x40EC - assert capa.features.String("cimage/png") in features # UTF-16 LE, offset 0x350E - - -def test_function_calls_to(sample_9324d1a8ae37a36ae560c37448c9705a): - features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60)) - assert capa.features.Characteristic("calls to") in features - assert len(features[capa.features.Characteristic("calls to")]) == 1 - - -def test_function_calls_to64(sample_lab21_01): - features = extract_function_features(viv_utils.Function(sample_lab21_01.vw, 0x1400052D0)) # memcpy - assert capa.features.Characteristic("calls to") in features - assert len(features[capa.features.Characteristic("calls to")]) == 8 - - -def test_function_calls_from(sample_9324d1a8ae37a36ae560c37448c9705a): - features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60)) - assert capa.features.Characteristic("calls from") in features - assert len(features[capa.features.Characteristic("calls from")]) == 23 - - -def test_basic_block_count(sample_9324d1a8ae37a36ae560c37448c9705a): - features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60)) - assert len(features[capa.features.basicblock.BasicBlock()]) == 26 - - -def test_indirect_call_features(sample_a933a1a402775cfa94b6bee0963f4b46): - features = extract_function_features(viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA68A0)) - assert capa.features.Characteristic("indirect call") in features - assert len(features[capa.features.Characteristic("indirect call")]) == 3 - - -def test_indirect_calls_resolved(sample_c91887d861d9bd4a5872249b641bc9f9): - features = extract_function_features(viv_utils.Function(sample_c91887d861d9bd4a5872249b641bc9f9.vw, 0x401A77)) - assert capa.features.insn.API("kernel32.CreatePipe") in features - assert capa.features.insn.API("kernel32.SetHandleInformation") in features - assert capa.features.insn.API("kernel32.CloseHandle") in features - assert capa.features.insn.API("kernel32.WriteFile") in features +@parametrize( + "sample,scope,feature,expected", FEATURE_COUNT_TESTS, indirect=["sample", "scope"], +) +def do_test_viv_feature_counts(sample, scope, feature, expected): + test_feature_count(get_viv_extractor, sample, scope, feature, expected)