From d23ef48bb6a426f0c5a7705e49bbbdaaae6b7231 Mon Sep 17 00:00:00 2001 From: William Ballenthin Date: Wed, 1 Jul 2020 12:33:13 -0600 Subject: [PATCH] pep8 --- tests/fixtures.py | 28 +- tests/test_engine.py | 361 +++++++++++++------- tests/test_fmt.py | 103 +++--- tests/test_freeze.py | 179 ++++++---- tests/test_helpers.py | 8 +- tests/test_main.py | 396 ++++++++++++--------- tests/test_rules.py | 679 +++++++++++++++++++++++-------------- tests/test_viv_features.py | 263 ++++++++------ 8 files changed, 1234 insertions(+), 783 deletions(-) diff --git a/tests/fixtures.py b/tests/fixtures.py index 5ba741ca..64862331 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -9,70 +9,74 @@ import viv_utils CD = os.path.dirname(__file__) -Sample = collections.namedtuple('Sample', ['vw', 'path']) +Sample = collections.namedtuple("Sample", ["vw", "path"]) @pytest.fixture def mimikatz(): - path = os.path.join(CD, 'data', 'mimikatz.exe_') + path = os.path.join(CD, "data", "mimikatz.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_a933a1a402775cfa94b6bee0963f4b46(): - path = os.path.join(CD, 'data', 'a933a1a402775cfa94b6bee0963f4b46.dll_') + path = os.path.join(CD, "data", "a933a1a402775cfa94b6bee0963f4b46.dll_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def kernel32(): - path = os.path.join(CD, 'data', 'kernel32.dll_') + path = os.path.join(CD, "data", "kernel32.dll_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_a198216798ca38f280dc413f8c57f2c2(): - path = os.path.join(CD, 'data', 'a198216798ca38f280dc413f8c57f2c2.exe_') + path = os.path.join(CD, "data", "a198216798ca38f280dc413f8c57f2c2.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_9324d1a8ae37a36ae560c37448c9705a(): - path = os.path.join(CD, 'data', '9324d1a8ae37a36ae560c37448c9705a.exe_') + path = os.path.join(CD, "data", "9324d1a8ae37a36ae560c37448c9705a.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def pma_lab_12_04(): - path = os.path.join(CD, 'data', 'Practical Malware Analysis Lab 12-04.exe_') + 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(): - path = os.path.join(CD, 'data', 'bfb9b5391a13d0afd787e87ab90f14f5.dll_') + path = os.path.join(CD, "data", "bfb9b5391a13d0afd787e87ab90f14f5.dll_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_lab21_01(): - path = os.path.join(CD, 'data', 'Practical Malware Analysis Lab 21-01.exe_') + 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(): - path = os.path.join(CD, 'data', 'c91887d861d9bd4a5872249b641bc9f9.exe_') + path = os.path.join(CD, "data", "c91887d861d9bd4a5872249b641bc9f9.exe_") return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41(): - path = os.path.join(CD, 'data', '39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.dll_') + path = os.path.join( + CD, + "data", + "39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.dll_", + ) return Sample(viv_utils.getWorkspace(path), path) @pytest.fixture def sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32(): - path = os.path.join(CD, 'data', '499c2a85f6e8142c3f48d4251c9c7cd6.raw32') + path = os.path.join(CD, "data", "499c2a85f6e8142c3f48d4251c9c7cd6.raw32") return Sample(viv_utils.getShellcodeWorkspace(path), path) diff --git a/tests/test_engine.py b/tests/test_engine.py index 5c7c9a3c..3f2ff7e0 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -41,25 +41,46 @@ def test_some(): assert Some(1, Number(1)).evaluate({Number(0): {1}}) == False assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}}) == False - assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}, Number(1): {1}}) == False - assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}}) == True - assert Some(2, Number(1), Number(2), Number(3)).evaluate( - {Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}}) == True - assert Some(2, Number(1), Number(2), Number(3)).evaluate( - {Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}, Number(4): {1}}) == True + assert ( + Some(2, Number(1), Number(2), Number(3)).evaluate( + {Number(0): {1}, Number(1): {1}} + ) + == False + ) + assert ( + Some(2, Number(1), Number(2), Number(3)).evaluate( + {Number(0): {1}, Number(1): {1}, Number(2): {1}} + ) + == True + ) + assert ( + Some(2, Number(1), Number(2), Number(3)).evaluate( + {Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}} + ) + == True + ) + assert ( + Some(2, Number(1), Number(2), Number(3)).evaluate( + { + Number(0): {1}, + Number(1): {1}, + Number(2): {1}, + Number(3): {1}, + Number(4): {1}, + } + ) + == True + ) def test_complex(): assert True == Or( And(Number(1), Number(2)), - Or(Number(3), - Some(2, Number(4), Number(5), Number(6))) + Or(Number(3), Some(2, Number(4), Number(5), Number(6))), ).evaluate({Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}}) assert False == Or( - And(Number(1), Number(2)), - Or(Number(3), - Some(2, Number(4), Number(5))) + And(Number(1), Number(2)), Or(Number(3), Some(2, Number(4), Number(5))) ).evaluate({Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}}) @@ -97,176 +118,252 @@ def test_range(): def test_match_adds_matched_rule_feature(): - '''show that using `match` adds a feature for matched rules.''' - rule = textwrap.dedent(''' + """show that using `match` adds a feature for matched rules.""" + rule = textwrap.dedent( + """ rule: meta: name: test rule features: - number: 100 - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) - features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1}}, 0x0) - assert capa.features.MatchedRule('test rule') in features + features, matches = capa.engine.match( + [r], {capa.features.insn.Number(100): {1}}, 0x0 + ) + assert capa.features.MatchedRule("test rule") in features def test_match_matched_rules(): - '''show that using `match` adds a feature for matched rules.''' + """show that using `match` adds a feature for matched rules.""" rules = [ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule1 - features: - - number: 100 - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule2 - features: - - match: test rule1 - ''')), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule1 + features: + - number: 100 + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule2 + features: + - match: test rule1 + """ + ) + ), ] - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.insn.Number(100): {1}}, 0x0) - assert capa.features.MatchedRule('test rule1') in features - assert capa.features.MatchedRule('test rule2') in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.insn.Number(100): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule1") in features + assert capa.features.MatchedRule("test rule2") in features # the ordering of the rules must not matter, # the engine should match rules in an appropriate order. - features, matches = capa.engine.match(capa.engine.topologically_order_rules(reversed(rules)), - {capa.features.insn.Number(100): {1}}, 0x0) - assert capa.features.MatchedRule('test rule1') in features - assert capa.features.MatchedRule('test rule2') in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(reversed(rules)), + {capa.features.insn.Number(100): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule1") in features + assert capa.features.MatchedRule("test rule2") in features def test_regex(): rules = [ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - and: - - string: /.*bbbb.*/ - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule with implied wildcards - features: - - and: - - string: /bbbb/ - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule with anchor - features: - - and: - - string: /^bbbb/ - ''')), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - and: + - string: /.*bbbb.*/ + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule with implied wildcards + features: + - and: + - string: /bbbb/ + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule with anchor + features: + - and: + - string: /^bbbb/ + """ + ) + ), ] - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.insn.Number(100): {1}}, 0x0) - assert capa.features.MatchedRule('test rule') not in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.insn.Number(100): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule") not in features - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.String('aaaa'): {1}}, 0x0) - assert capa.features.MatchedRule('test rule') not in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.String("aaaa"): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule") not in features - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.String('aBBBBa'): {1}}, 0x0) - assert capa.features.MatchedRule('test rule') not in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.String("aBBBBa"): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule") not in features - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.String('abbbba'): {1}}, 0x0) - assert capa.features.MatchedRule('test rule') in features - assert capa.features.MatchedRule('rule with implied wildcards') in features - assert capa.features.MatchedRule('rule with anchor') not in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.String("abbbba"): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule") in features + assert capa.features.MatchedRule("rule with implied wildcards") in features + assert capa.features.MatchedRule("rule with anchor") not in features def test_regex_ignorecase(): rules = [ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - and: - - string: /.*bbbb.*/i - ''')), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - and: + - string: /.*bbbb.*/i + """ + ) + ), ] - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.String('aBBBBa'): {1}}, 0x0) - assert capa.features.MatchedRule('test rule') in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.String("aBBBBa"): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule") in features def test_regex_complex(): rules = [ - capa.rules.Rule.from_yaml(textwrap.dedent(r''' - rule: - meta: - name: test rule - features: - - or: - - string: /.*HARDWARE\\Key\\key with spaces\\.*/i - ''')), + capa.rules.Rule.from_yaml( + textwrap.dedent( + r""" + rule: + meta: + name: test rule + features: + - or: + - string: /.*HARDWARE\\Key\\key with spaces\\.*/i + """ + ) + ), ] - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.String(r'Hardware\Key\key with spaces\some value'): {1}}, 0x0) - assert capa.features.MatchedRule('test rule') in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.String(r"Hardware\Key\key with spaces\some value"): {1}}, + 0x0, + ) + assert capa.features.MatchedRule("test rule") in features def test_match_namespace(): rules = [ - capa.rules.Rule.from_yaml(textwrap.dedent(''' + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ rule: meta: name: CreateFile API namespace: file/create/CreateFile features: - api: CreateFile - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ rule: meta: name: WriteFile API namespace: file/write features: - api: WriteFile - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: file-create - features: - - match: file/create - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: filesystem-any - features: - - match: file - ''')), + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: file-create + features: + - match: file/create + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: filesystem-any + features: + - match: file + """ + ) + ), ] - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.insn.API('CreateFile'): {1}}, - 0x0) - assert 'CreateFile API' in matches - assert 'file-create' in matches - assert 'filesystem-any' in matches - assert capa.features.MatchedRule('file') in features - assert capa.features.MatchedRule('file/create') in features - assert capa.features.MatchedRule('file/create/CreateFile') in features + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.insn.API("CreateFile"): {1}}, + 0x0, + ) + assert "CreateFile API" in matches + assert "file-create" in matches + assert "filesystem-any" in matches + assert capa.features.MatchedRule("file") in features + assert capa.features.MatchedRule("file/create") in features + assert capa.features.MatchedRule("file/create/CreateFile") in features - features, matches = capa.engine.match(capa.engine.topologically_order_rules(rules), - {capa.features.insn.API('WriteFile'): {1}}, - 0x0) - assert 'WriteFile API' in matches - assert 'file-create' not in matches - assert 'filesystem-any' in matches + features, matches = capa.engine.match( + capa.engine.topologically_order_rules(rules), + {capa.features.insn.API("WriteFile"): {1}}, + 0x0, + ) + assert "WriteFile API" in matches + assert "file-create" not in matches + assert "filesystem-any" in matches diff --git a/tests/test_fmt.py b/tests/test_fmt.py index 3c6ddc14..33fe077c 100644 --- a/tests/test_fmt.py +++ b/tests/test_fmt.py @@ -2,24 +2,27 @@ import textwrap import capa.rules -EXPECTED = textwrap.dedent('''\ - rule: - meta: - name: test rule - author: user@domain.com - scope: function - examples: - - foo1234 - - bar5678 - features: - - and: - - number: 1 - - number: 2 -''') +EXPECTED = textwrap.dedent( + """\ + rule: + meta: + name: test rule + author: user@domain.com + scope: function + examples: + - foo1234 + - bar5678 + features: + - and: + - number: 1 + - number: 2 + """ +) def test_rule_reformat_top_level_elements(): - rule = textwrap.dedent('''\ + rule = textwrap.dedent( + """ rule: features: - and: @@ -31,13 +34,16 @@ def test_rule_reformat_top_level_elements(): scope: function examples: - foo1234 - - bar5678''') + - bar5678 + """ + ) assert capa.rules.Rule.from_yaml(rule).to_yaml() == EXPECTED def test_rule_reformat_indentation(): - rule = textwrap.dedent('''\ + rule = textwrap.dedent( + """ rule: meta: name: test rule @@ -49,45 +55,52 @@ def test_rule_reformat_indentation(): features: - and: - number: 1 - - number: 2''') + - number: 2 + """ + ) assert capa.rules.Rule.from_yaml(rule).to_yaml() == EXPECTED def test_rule_reformat_order(): - rule = textwrap.dedent('''\ - rule: - meta: - author: user@domain.com - examples: - - foo1234 - - bar5678 - scope: function - name: test rule - features: - - and: - - number: 1 - - number: 2''') + rule = textwrap.dedent( + """ + rule: + meta: + author: user@domain.com + examples: + - foo1234 + - bar5678 + scope: function + name: test rule + features: + - and: + - number: 1 + - number: 2 + """ + ) assert capa.rules.Rule.from_yaml(rule).to_yaml() == EXPECTED def test_rule_reformat_meta_update(): - rule = textwrap.dedent('''\ - rule: - meta: - author: user@domain.com - examples: - - foo1234 - - bar5678 - scope: function - name: AAAA - features: - - and: - - number: 1 - - number: 2''') + rule = textwrap.dedent( + """ + rule: + meta: + author: user@domain.com + examples: + - foo1234 + - bar5678 + scope: function + name: AAAA + features: + - and: + - number: 1 + - number: 2 + """ + ) rule = capa.rules.Rule.from_yaml(rule) rule.name = "test rule" assert rule.to_yaml() == EXPECTED - diff --git a/tests/test_freeze.py b/tests/test_freeze.py index 29c05fde..a78333b9 100644 --- a/tests/test_freeze.py +++ b/tests/test_freeze.py @@ -10,38 +10,44 @@ import capa.features.freeze from fixtures import * -EXTRACTOR = capa.features.extractors.NullFeatureExtractor({ - 'file features': [ - (0x402345, capa.features.Characteristic('embedded pe', True)), - ], - 'functions': { - 0x401000: { - 'features': [ - (0x401000, capa.features.Characteristic('switch', True)), - ], - 'basic blocks': { - 0x401000: { - 'features': [ - (0x401000, capa.features.Characteristic('tight loop', True)), - ], - 'instructions': { - 0x401000: { - 'features': [ - (0x401000, capa.features.insn.Mnemonic('xor')), - (0x401000, capa.features.Characteristic('nzxor', True)), - ], +EXTRACTOR = capa.features.extractors.NullFeatureExtractor( + { + "file features": [ + (0x402345, capa.features.Characteristic("embedded pe", True)), + ], + "functions": { + 0x401000: { + "features": [(0x401000, capa.features.Characteristic("switch", True)),], + "basic blocks": { + 0x401000: { + "features": [ + ( + 0x401000, + capa.features.Characteristic("tight loop", True), + ), + ], + "instructions": { + 0x401000: { + "features": [ + (0x401000, capa.features.insn.Mnemonic("xor")), + ( + 0x401000, + capa.features.Characteristic("nzxor", True), + ), + ], + }, + 0x401002: { + "features": [ + (0x401002, capa.features.insn.Mnemonic("mov")), + ] + }, }, - 0x401002: { - 'features': [ - (0x401002, capa.features.insn.Mnemonic('mov')), - ] - } - } + }, }, - } + }, }, } -}) +) def test_null_feature_extractor(): @@ -49,29 +55,35 @@ def test_null_feature_extractor(): assert list(EXTRACTOR.get_basic_blocks(0x401000)) == [0x401000] assert list(EXTRACTOR.get_instructions(0x401000, 0x0401000)) == [0x401000, 0x401002] - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: xor loop - scope: basic block - features: - - and: - - characteristic(tight loop): true - - mnemonic: xor - - characteristic(nzxor): true - ''')), - ]) + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: xor loop + scope: basic block + features: + - and: + - characteristic(tight loop): true + - mnemonic: xor + - characteristic(nzxor): true + """ + ) + ), + ] + ) capabilities = capa.main.find_capabilities(rules, EXTRACTOR) - assert 'xor loop' in capabilities + assert "xor loop" in capabilities def compare_extractors(a, b): - ''' + """ args: a (capa.features.extractors.NullFeatureExtractor) b (capa.features.extractors.NullFeatureExtractor) - ''' + """ # TODO: ordering of these things probably doesn't work yet @@ -79,44 +91,62 @@ def compare_extractors(a, b): assert list(a.get_functions()) == list(b.get_functions()) for f in a.get_functions(): assert list(a.get_basic_blocks(f)) == list(b.get_basic_blocks(f)) - assert list(a.extract_function_features(f)) == list(b.extract_function_features(f)) + assert list(a.extract_function_features(f)) == list( + b.extract_function_features(f) + ) for bb in a.get_basic_blocks(f): assert list(a.get_instructions(f, bb)) == list(b.get_instructions(f, bb)) - assert list(a.extract_basic_block_features(f, bb)) == list(b.extract_basic_block_features(f, bb)) + assert list(a.extract_basic_block_features(f, bb)) == list( + b.extract_basic_block_features(f, bb) + ) for insn in a.get_instructions(f, bb): - assert list(a.extract_insn_features(f, bb, insn)) == list(b.extract_insn_features(f, bb, insn)) + assert list(a.extract_insn_features(f, bb, insn)) == list( + b.extract_insn_features(f, bb, insn) + ) def compare_extractors_viv_null(viv_ext, null_ext): - ''' + """ almost identical to compare_extractors but adds casts to ints since the VivisectFeatureExtractor returns objects and NullFeatureExtractor returns ints args: viv_ext (capa.features.extractors.viv.VivisectFeatureExtractor) null_ext (capa.features.extractors.NullFeatureExtractor) - ''' + """ # TODO: ordering of these things probably doesn't work yet - assert list(viv_ext.extract_file_features()) == list(null_ext.extract_file_features()) + assert list(viv_ext.extract_file_features()) == list( + null_ext.extract_file_features() + ) assert to_int(list(viv_ext.get_functions())) == list(null_ext.get_functions()) for f in viv_ext.get_functions(): - assert to_int(list(viv_ext.get_basic_blocks(f))) == list(null_ext.get_basic_blocks(to_int(f))) - assert list(viv_ext.extract_function_features(f)) == list(null_ext.extract_function_features(to_int(f))) + assert to_int(list(viv_ext.get_basic_blocks(f))) == list( + null_ext.get_basic_blocks(to_int(f)) + ) + assert list(viv_ext.extract_function_features(f)) == list( + null_ext.extract_function_features(to_int(f)) + ) for bb in viv_ext.get_basic_blocks(f): - assert to_int(list(viv_ext.get_instructions(f, bb))) == list(null_ext.get_instructions(to_int(f), to_int(bb))) - assert list(viv_ext.extract_basic_block_features(f, bb)) == list(null_ext.extract_basic_block_features(to_int(f), to_int(bb))) + assert to_int(list(viv_ext.get_instructions(f, bb))) == list( + null_ext.get_instructions(to_int(f), to_int(bb)) + ) + assert list(viv_ext.extract_basic_block_features(f, bb)) == list( + null_ext.extract_basic_block_features(to_int(f), to_int(bb)) + ) for insn in viv_ext.get_instructions(f, bb): - assert list(viv_ext.extract_insn_features(f, bb, insn)) == list(null_ext.extract_insn_features(to_int(f), to_int(bb), to_int(insn))) + assert list(viv_ext.extract_insn_features(f, bb, insn)) == list( + null_ext.extract_insn_features(to_int(f), to_int(bb), to_int(insn)) + ) def to_int(o): - '''helper to get int value of extractor items''' + """helper to get int value of extractor items""" if isinstance(o, list): return map(lambda x: capa.helpers.oint(x), o) else: @@ -144,30 +174,37 @@ def roundtrip_feature(feature): def test_serialize_features(): - roundtrip_feature(capa.features.insn.API('advapi32.CryptAcquireContextW')) - roundtrip_feature(capa.features.String('SCardControl')) + roundtrip_feature(capa.features.insn.API("advapi32.CryptAcquireContextW")) + roundtrip_feature(capa.features.String("SCardControl")) roundtrip_feature(capa.features.insn.Number(0xFF)) roundtrip_feature(capa.features.insn.Offset(0x0)) - roundtrip_feature(capa.features.insn.Mnemonic('push')) - roundtrip_feature(capa.features.file.Section('.rsrc')) - roundtrip_feature(capa.features.Characteristic('tight loop', True)) + roundtrip_feature(capa.features.insn.Mnemonic("push")) + roundtrip_feature(capa.features.file.Section(".rsrc")) + roundtrip_feature(capa.features.Characteristic("tight loop", True)) roundtrip_feature(capa.features.basicblock.BasicBlock()) - roundtrip_feature(capa.features.file.Export('BaseThreadInitThunk')) - roundtrip_feature(capa.features.file.Import('kernel32.IsWow64Process')) - roundtrip_feature(capa.features.file.Import('#11')) + roundtrip_feature(capa.features.file.Export("BaseThreadInitThunk")) + roundtrip_feature(capa.features.file.Import("kernel32.IsWow64Process")) + roundtrip_feature(capa.features.file.Import("#11")) def test_freeze_sample(tmpdir, sample_9324d1a8ae37a36ae560c37448c9705a): # tmpdir fixture handles cleanup - o = tmpdir.mkdir('capa').join('test.frz').strpath - assert capa.features.freeze.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, o, '-v']) == 0 + o = tmpdir.mkdir("capa").join("test.frz").strpath + assert ( + capa.features.freeze.main( + [sample_9324d1a8ae37a36ae560c37448c9705a.path, o, "-v"] + ) + == 0 + ) def test_freeze_load_sample(tmpdir, sample_9324d1a8ae37a36ae560c37448c9705a): - o = tmpdir.mkdir('capa').join('test.frz') - viv_extractor = capa.features.extractors.viv.VivisectFeatureExtractor(sample_9324d1a8ae37a36ae560c37448c9705a.vw, - sample_9324d1a8ae37a36ae560c37448c9705a.path) - with open(o.strpath, 'wb') as f: + o = tmpdir.mkdir("capa").join("test.frz") + viv_extractor = capa.features.extractors.viv.VivisectFeatureExtractor( + sample_9324d1a8ae37a36ae560c37448c9705a.vw, + sample_9324d1a8ae37a36ae560c37448c9705a.path, + ) + with open(o.strpath, "wb") as f: f.write(capa.features.freeze.dump(viv_extractor)) - null_extractor = capa.features.freeze.load(o.open('rb').read()) + null_extractor = capa.features.freeze.load(o.open("rb").read()) compare_extractors_viv_null(viv_extractor, null_extractor) diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 731e6691..c3dbe5e3 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -6,10 +6,10 @@ from capa.features.extractors import helpers def test_all_zeros(): # Python 2: # Python 3: - a = b'\x00\x00\x00\x00' - b = codecs.decode('00000000', 'hex') - c = b'\x01\x00\x00\x00' - d = codecs.decode('01000000', 'hex') + a = b"\x00\x00\x00\x00" + b = codecs.decode("00000000", "hex") + c = b"\x01\x00\x00\x00" + d = codecs.decode("01000000", "hex") assert helpers.all_zeros(a) is True assert helpers.all_zeros(b) is True assert helpers.all_zeros(c) is False diff --git a/tests/test_main.py b/tests/test_main.py index 5c5640bc..8640c4f1 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -12,192 +12,274 @@ from fixtures import * def test_main(sample_9324d1a8ae37a36ae560c37448c9705a): # tests rules can be loaded successfully - assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, '-v']) == 0 + assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, "-v"]) == 0 def test_main_single_rule(sample_9324d1a8ae37a36ae560c37448c9705a, tmpdir): # tests a single rule can be loaded successfully - RULE_CONTENT = textwrap.dedent(''' - rule: - meta: - name: test rule - scope: file - features: - - string: test - ''') - rule_file = tmpdir.mkdir('capa').join('rule.yml') + RULE_CONTENT = textwrap.dedent( + """ + rule: + meta: + name: test rule + scope: file + features: + - string: test + """ + ) + rule_file = tmpdir.mkdir("capa").join("rule.yml") rule_file.write(RULE_CONTENT) - assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, '-v', '-r', rule_file.strpath]) == 0 + assert ( + capa.main.main( + [ + sample_9324d1a8ae37a36ae560c37448c9705a.path, + "-v", + "-r", + rule_file.strpath, + ] + ) + == 0 + ) def test_main_shellcode(sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32): - assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, '-v', '-f', 'sc32']) == 0 + assert ( + capa.main.main( + [sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, "-v", "-f", "sc32"] + ) + == 0 + ) def test_ruleset(): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: file rule - scope: file - features: - - characteristic(embedded pe): y - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: function rule - scope: function - features: - - characteristic(switch): y - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: basic block rule - scope: basic block - features: - - characteristic(nzxor): y - ''')), - - ]) + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: file rule + scope: file + features: + - characteristic(embedded pe): y + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: function rule + scope: function + features: + - characteristic(switch): y + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: basic block rule + scope: basic block + features: + - characteristic(nzxor): y + """ + ) + ), + ] + ) assert len(rules.file_rules) == 1 assert len(rules.function_rules) == 1 assert len(rules.basic_block_rules) == 1 def test_match_across_scopes_file_function(sample_9324d1a8ae37a36ae560c37448c9705a): - rules = capa.rules.RuleSet([ - # this rule should match on a function (0x4073F0) - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: install service - scope: function - examples: - - 9324d1a8ae37a36ae560c37448c9705a:0x4073F0 - features: - - and: - - api: advapi32.OpenSCManagerA - - api: advapi32.CreateServiceA - - api: advapi32.StartServiceA - ''')), - # this rule should match on a file feature - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: .text section - scope: file - examples: - - 9324d1a8ae37a36ae560c37448c9705a - features: - - section: .text - ''')), - # this rule should match on earlier rule matches: - # - install service, with function scope - # - .text section, with file scope - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: .text section and install service - scope: file - examples: - - 9324d1a8ae37a36ae560c37448c9705a - features: - - and: - - match: install service - - match: .text section - ''')), - ]) - extractor = capa.features.extractors.viv.VivisectFeatureExtractor(sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path) + rules = capa.rules.RuleSet( + [ + # this rule should match on a function (0x4073F0) + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: install service + scope: function + examples: + - 9324d1a8ae37a36ae560c37448c9705a:0x4073F0 + features: + - and: + - api: advapi32.OpenSCManagerA + - api: advapi32.CreateServiceA + - api: advapi32.StartServiceA + """ + ) + ), + # this rule should match on a file feature + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: .text section + scope: file + examples: + - 9324d1a8ae37a36ae560c37448c9705a + features: + - section: .text + """ + ) + ), + # this rule should match on earlier rule matches: + # - install service, with function scope + # - .text section, with file scope + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: .text section and install service + scope: file + examples: + - 9324d1a8ae37a36ae560c37448c9705a + features: + - and: + - match: install service + - match: .text section + """ + ) + ), + ] + ) + extractor = capa.features.extractors.viv.VivisectFeatureExtractor( + sample_9324d1a8ae37a36ae560c37448c9705a.vw, + sample_9324d1a8ae37a36ae560c37448c9705a.path, + ) capabilities = capa.main.find_capabilities(rules, extractor) - assert 'install service' in capabilities - assert '.text section' in capabilities - assert '.text section and install service' in capabilities + assert "install service" in capabilities + assert ".text section" in capabilities + assert ".text section and install service" in capabilities def test_match_across_scopes(sample_9324d1a8ae37a36ae560c37448c9705a): - rules = capa.rules.RuleSet([ - # this rule should match on a basic block (including at least 0x403685) - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: tight loop - scope: basic block - examples: - - 9324d1a8ae37a36ae560c37448c9705a:0x403685 - features: - - characteristic(tight loop): true - ''')), - # this rule should match on a function (0x403660) - # based on API, as well as prior basic block rule match - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: kill thread loop - scope: function - examples: - - 9324d1a8ae37a36ae560c37448c9705a:0x403660 - features: - - and: - - api: kernel32.TerminateThread - - api: kernel32.CloseHandle - - match: tight loop - ''')), - # this rule should match on a file feature and a prior function rule match - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: kill thread program - scope: file - examples: - - 9324d1a8ae37a36ae560c37448c9705a - features: - - and: - - section: .text - - match: kill thread loop - ''')), - ]) - extractor = capa.features.extractors.viv.VivisectFeatureExtractor(sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path) + rules = capa.rules.RuleSet( + [ + # this rule should match on a basic block (including at least 0x403685) + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: tight loop + scope: basic block + examples: + - 9324d1a8ae37a36ae560c37448c9705a:0x403685 + features: + - characteristic(tight loop): true + """ + ) + ), + # this rule should match on a function (0x403660) + # based on API, as well as prior basic block rule match + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: kill thread loop + scope: function + examples: + - 9324d1a8ae37a36ae560c37448c9705a:0x403660 + features: + - and: + - api: kernel32.TerminateThread + - api: kernel32.CloseHandle + - match: tight loop + """ + ) + ), + # this rule should match on a file feature and a prior function rule match + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: kill thread program + scope: file + examples: + - 9324d1a8ae37a36ae560c37448c9705a + features: + - and: + - section: .text + - match: kill thread loop + """ + ) + ), + ] + ) + extractor = capa.features.extractors.viv.VivisectFeatureExtractor( + sample_9324d1a8ae37a36ae560c37448c9705a.vw, + sample_9324d1a8ae37a36ae560c37448c9705a.path, + ) capabilities = capa.main.find_capabilities(rules, extractor) - assert 'tight loop' in capabilities - assert 'kill thread loop' in capabilities - assert 'kill thread program' in capabilities + assert "tight loop" in capabilities + assert "kill thread loop" in capabilities + assert "kill thread program" in capabilities def test_subscope_bb_rules(sample_9324d1a8ae37a36ae560c37448c9705a): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - scope: function - features: - - and: - - basic block: - - characteristic(tight loop): true - ''')) - ]) + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + scope: function + features: + - and: + - basic block: + - characteristic(tight loop): true + """ + ) + ) + ] + ) # tight loop at 0x403685 - extractor = capa.features.extractors.viv.VivisectFeatureExtractor(sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path) + extractor = capa.features.extractors.viv.VivisectFeatureExtractor( + sample_9324d1a8ae37a36ae560c37448c9705a.vw, + sample_9324d1a8ae37a36ae560c37448c9705a.path, + ) capabilities = capa.main.find_capabilities(rules, extractor) - assert 'test rule' in capabilities + assert "test rule" in capabilities def test_byte_matching(sample_9324d1a8ae37a36ae560c37448c9705a): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: byte match test - scope: function - features: - - and: - - bytes: ED 24 9E F4 52 A9 07 47 55 8E E1 AB 30 8E 23 61 - ''')) - ]) + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: byte match test + scope: function + features: + - and: + - bytes: ED 24 9E F4 52 A9 07 47 55 8E E1 AB 30 8E 23 61 + """ + ) + ) + ] + ) - extractor = capa.features.extractors.viv.VivisectFeatureExtractor(sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path) + extractor = capa.features.extractors.viv.VivisectFeatureExtractor( + sample_9324d1a8ae37a36ae560c37448c9705a.vw, + sample_9324d1a8ae37a36ae560c37448c9705a.path, + ) capabilities = capa.main.find_capabilities(rules, extractor) - assert 'byte match test' in capabilities + assert "byte match test" in capabilities diff --git a/tests/test_rules.py b/tests/test_rules.py index 98bd2e79..95b88d24 100644 --- a/tests/test_rules.py +++ b/tests/test_rules.py @@ -7,13 +7,14 @@ from capa.features.insn import Number, Offset def test_rule_ctor(): - r = capa.rules.Rule('test rule', capa.rules.FUNCTION_SCOPE, Number(1), {}) + r = capa.rules.Rule("test rule", capa.rules.FUNCTION_SCOPE, Number(1), {}) assert r.evaluate({Number(0): {1}}) == False assert r.evaluate({Number(1): {1}}) == True def test_rule_yaml(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule @@ -26,16 +27,21 @@ def test_rule_yaml(): - and: - number: 1 - number: 2 - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) assert r.evaluate({Number(0): {1}}) == False assert r.evaluate({Number(0): {1}, Number(1): {1}}) == False assert r.evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}}) == True - assert r.evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}}) == True + assert ( + r.evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}}) + == True + ) def test_rule_yaml_complex(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule @@ -50,36 +56,44 @@ def test_rule_yaml_complex(): - number: 4 - number: 5 - number: 6 - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) - assert r.evaluate({Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}}) == True + assert ( + r.evaluate({Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}}) + == True + ) assert r.evaluate({Number(6): {1}, Number(7): {1}, Number(8): {1}}) == False def test_rule_yaml_not(): - rule = textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - and: - - number: 1 - - not: - - number: 2 - ''') + rule = textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - and: + - number: 1 + - not: + - number: 2 + """ + ) r = capa.rules.Rule.from_yaml(rule) assert r.evaluate({Number(1): {1}}) == True assert r.evaluate({Number(1): {1}, Number(2): {1}}) == False def test_rule_yaml_count(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule features: - count(number(100)): 1 - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) assert r.evaluate({Number(100): {}}) == False assert r.evaluate({Number(100): {1}}) == True @@ -87,13 +101,15 @@ def test_rule_yaml_count(): def test_rule_yaml_count_range(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule features: - count(number(100)): (1, 2) - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) assert r.evaluate({Number(100): {}}) == False assert r.evaluate({Number(100): {1}}) == True @@ -103,83 +119,115 @@ def test_rule_yaml_count_range(): def test_invalid_rule_feature(): with pytest.raises(capa.rules.InvalidRule): - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - foo: true - ''')) + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - foo: true + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - scope: file - features: - - characteristic(nzxor): true - ''')) + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + scope: file + features: + - characteristic(nzxor): true + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - scope: function - features: - - characteristic(embedded pe): true - ''')) + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + scope: function + features: + - characteristic(embedded pe): true + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - scope: basic block - features: - - characteristic(embedded pe): true - ''')) + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + scope: basic block + features: + - characteristic(embedded pe): true + """ + ) + ) def test_lib_rules(): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: a lib rule - lib: true - features: - - api: CreateFileA - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: a standard rule - lib: false - features: - - api: CreateFileW - ''')), - ]) + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: a lib rule + lib: true + features: + - api: CreateFileA + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: a standard rule + lib: false + features: + - api: CreateFileW + """ + ) + ), + ] + ) assert len(rules.function_rules) == 1 def test_subscope_rules(): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - scope: file - features: - - and: - - characteristic(embedded pe): true - - function: + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + scope: file + features: - and: - - characteristic(nzxor): true - - characteristic(switch): true - ''')) - ]) + - characteristic(embedded pe): true + - function: + - and: + - characteristic(nzxor): true + - characteristic(switch): true + """ + ) + ) + ] + ) # the file rule scope will have one rules: # - `test rule` assert len(rules.file_rules) == 1 @@ -191,59 +239,84 @@ def test_subscope_rules(): def test_duplicate_rules(): with pytest.raises(capa.rules.InvalidRule): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule-name - features: - - api: CreateFileA - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule-name - features: - - api: CreateFileW - ''')), - ]) + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule-name + features: + - api: CreateFileA + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule-name + features: + - api: CreateFileW + """ + ) + ), + ] + ) def test_missing_dependency(): with pytest.raises(capa.rules.InvalidRule): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: dependent rule - features: - - match: missing rule - ''')), - ]) + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: dependent rule + features: + - match: missing rule + """ + ) + ), + ] + ) def test_invalid_rules(): with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ rule: meta: name: test rule features: - characteristic(number(1)): True - ''')) + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ rule: meta: name: test rule features: - characteristic(count(number(100))): True - ''')) + """ + ) + ) def test_number_symbol(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule @@ -256,19 +329,21 @@ def test_number_symbol(): - number: 4 = symbol name = another name - number: 0x100 = symbol name - number: 0x11 = (FLAG_A | FLAG_B) - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) children = list(r.statement.get_children()) assert (Number(1) in children) == True assert (Number(0xFFFFFFFF) in children) == True - assert (Number(2, 'symbol name') in children) == True - assert (Number(3, 'symbol name') in children) == True - assert (Number(4, 'symbol name = another name') in children) == True - assert (Number(0x100, 'symbol name') in children) == True + assert (Number(2, "symbol name") in children) == True + assert (Number(3, "symbol name") in children) == True + assert (Number(4, "symbol name = another name") in children) == True + assert (Number(0x100, "symbol name") in children) == True def test_count_number_symbol(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule @@ -277,46 +352,60 @@ def test_count_number_symbol(): - count(number(2 = symbol name)): 1 - count(number(0x100 = symbol name)): 2 or more - count(number(0x11 = (FLAG_A | FLAG_B))): 2 or more - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) assert r.evaluate({Number(2): {}}) == False assert r.evaluate({Number(2): {1}}) == True assert r.evaluate({Number(2): {1, 2}}) == False - assert r.evaluate({Number(0x100, 'symbol name'): {1}}) == False - assert r.evaluate({Number(0x100, 'symbol name'): {1, 2, 3}}) == True + assert r.evaluate({Number(0x100, "symbol name"): {1}}) == False + assert r.evaluate({Number(0x100, "symbol name"): {1, 2, 3}}) == True def test_invalid_number(): with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - number: "this is a string" - ''')) + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - number: "this is a string" + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - number: 2= - ''')) + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - number: 2= + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - number: symbol name = 2 - ''')) + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - number: symbol name = 2 + """ + ) + ) def test_offset_symbol(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule @@ -327,18 +416,20 @@ def test_offset_symbol(): - offset: 3 = symbol name - offset: 4 = symbol name = another name - offset: 0x100 = symbol name - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) children = list(r.statement.get_children()) assert (Offset(1) in children) == True - assert (Offset(2, 'symbol name') in children) == True - assert (Offset(3, 'symbol name') in children) == True - assert (Offset(4, 'symbol name = another name') in children) == True - assert (Offset(0x100, 'symbol name') in children) == True + assert (Offset(2, "symbol name") in children) == True + assert (Offset(3, "symbol name") in children) == True + assert (Offset(4, "symbol name = another name") in children) == True + assert (Offset(0x100, "symbol name") in children) == True def test_count_offset_symbol(): - rule = textwrap.dedent(''' + rule = textwrap.dedent( + """ rule: meta: name: test rule @@ -347,153 +438,215 @@ def test_count_offset_symbol(): - count(offset(2 = symbol name)): 1 - count(offset(0x100 = symbol name)): 2 or more - count(offset(0x11 = (FLAG_A | FLAG_B))): 2 or more - ''') + """ + ) r = capa.rules.Rule.from_yaml(rule) assert r.evaluate({Offset(2): {}}) == False assert r.evaluate({Offset(2): {1}}) == True assert r.evaluate({Offset(2): {1, 2}}) == False - assert r.evaluate({Offset(0x100, 'symbol name'): {1}}) == False - assert r.evaluate({Offset(0x100, 'symbol name'): {1, 2, 3}}) == True + assert r.evaluate({Offset(0x100, "symbol name"): {1}}) == False + assert r.evaluate({Offset(0x100, "symbol name"): {1, 2, 3}}) == True def test_invalid_offset(): with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - offset: "this is a string" - ''')) + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - offset: "this is a string" + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - offset: 2= - ''')) + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - offset: 2= + """ + ) + ) with pytest.raises(capa.rules.InvalidRule): - r = capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: test rule - features: - - offset: symbol name = 2 - ''')) + r = capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: test rule + features: + - offset: symbol name = 2 + """ + ) + ) def test_filter_rules(): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule 1 - author: joe - features: - - api: CreateFile - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule 2 - features: - - string: joe - ''')), - ]) - rules = rules.filter_rules_by_meta('joe') - assert len(rules) == 1 - assert ('rule 1' in rules.rules) - - -def test_filter_rules_dependencies(): - rules = capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule 1 - features: - - match: rule 2 - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule 2 - features: - - match: rule 3 - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule 3 - features: - - api: CreateFile - ''')), - ]) - rules = rules.filter_rules_by_meta('rule 1') - assert(len(rules.rules) == 3) - assert('rule 1' in rules.rules) - assert('rule 2' in rules.rules) - assert('rule 3' in rules.rules) - - -def test_filter_rules_missing_dependency(): - with pytest.raises(capa.rules.InvalidRule): - capa.rules.RuleSet([ - capa.rules.Rule.from_yaml(textwrap.dedent(''' + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ rule: meta: name: rule 1 author: joe + features: + - api: CreateFile + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule 2 + features: + - string: joe + """ + ) + ), + ] + ) + rules = rules.filter_rules_by_meta("joe") + assert len(rules) == 1 + assert "rule 1" in rules.rules + + +def test_filter_rules_dependencies(): + rules = capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule 1 features: - match: rule 2 - ''')), - ]) + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule 2 + features: + - match: rule 3 + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule 3 + features: + - api: CreateFile + """ + ) + ), + ] + ) + rules = rules.filter_rules_by_meta("rule 1") + assert len(rules.rules) == 3 + assert "rule 1" in rules.rules + assert "rule 2" in rules.rules + assert "rule 3" in rules.rules + + +def test_filter_rules_missing_dependency(): + with pytest.raises(capa.rules.InvalidRule): + capa.rules.RuleSet( + [ + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule 1 + author: joe + features: + - match: rule 2 + """ + ) + ), + ] + ) def test_rules_namespace_dependencies(): rules = [ - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule 1 - namespace: ns1/nsA - features: - - api: CreateFile - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' - rule: - meta: - name: rule 2 - namespace: ns1/nsB - features: - - api: CreateFile - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule 1 + namespace: ns1/nsA + features: + - api: CreateFile + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ + rule: + meta: + name: rule 2 + namespace: ns1/nsB + features: + - api: CreateFile + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ rule: meta: name: rule 3 features: - match: ns1/nsA - ''')), - capa.rules.Rule.from_yaml(textwrap.dedent(''' + """ + ) + ), + capa.rules.Rule.from_yaml( + textwrap.dedent( + """ rule: meta: name: rule 4 features: - match: ns1 - ''')), + """ + ) + ), ] - r3 = set(map(lambda r: r.name, capa.rules.get_rules_and_dependencies(rules, 'rule 3'))) - assert 'rule 1' in r3 - assert 'rule 2' not in r3 - assert 'rule 4' not in r3 - - r4 = set(map(lambda r: r.name, capa.rules.get_rules_and_dependencies(rules, 'rule 4'))) - assert 'rule 1' in r4 - assert 'rule 2' in r4 - assert 'rule 3' not in r4 + r3 = set( + map(lambda r: r.name, capa.rules.get_rules_and_dependencies(rules, "rule 3")) + ) + assert "rule 1" in r3 + assert "rule 2" not in r3 + assert "rule 4" not in r3 + r4 = set( + map(lambda r: r.name, capa.rules.get_rules_and_dependencies(rules, "rule 4")) + ) + assert "rule 1" in r4 + assert "rule 2" in r4 + assert "rule 3" not in r4 diff --git a/tests/test_viv_features.py b/tests/test_viv_features.py index ac0bac9d..c32606aa 100644 --- a/tests/test_viv_features.py +++ b/tests/test_viv_features.py @@ -1,5 +1,3 @@ -import collections - import viv_utils import capa.features @@ -26,9 +24,13 @@ 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): + 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): + 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) @@ -38,7 +40,9 @@ def extract_function_features(f): 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): + 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) @@ -47,50 +51,60 @@ def extract_basic_block_features(f, bb): 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 + 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 + 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 + 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 + 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 + assert capa.features.String("bcrypt.dll") not 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') + 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') + 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 @@ -116,11 +130,11 @@ def test_offset_features(mimikatz): def test_nzxor_features(mimikatz): features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x410DFC)) - assert capa.features.Characteristic('nzxor', True) in features # 0x0410F0B + assert capa.features.Characteristic("nzxor", True) in features # 0x0410F0B def get_bb_insn(f, va): - '''fetch the BasicBlock and Instruction instances for the given VA in the given function.''' + """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: @@ -133,7 +147,9 @@ def test_is_security_cookie(mimikatz): 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 + 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) @@ -144,24 +160,26 @@ def test_is_security_cookie(mimikatz): 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("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 + 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', True) in features + features = extract_function_features( + viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA6FEC) + ) + assert capa.features.Characteristic("peb access", True) 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 + 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): @@ -170,7 +188,7 @@ def test_tight_loop_features(mimikatz): if bb.va != 0x402F8E: continue features = extract_basic_block_features(f, bb) - assert capa.features.Characteristic('tight loop', True) in features + assert capa.features.Characteristic("tight loop", True) in features assert capa.features.basicblock.BasicBlock() in features @@ -180,119 +198,166 @@ def test_tight_loop_bb_features(mimikatz): if bb.va != 0x402F8E: continue features = extract_basic_block_features(f, bb) - assert capa.features.Characteristic('tight loop', True) in features + assert capa.features.Characteristic("tight loop", True) 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 + 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 + 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', True) in features + features = extract_function_features( + viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x4014D0) + ) + assert capa.features.Characteristic("cross section flow", True) 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', True) not in features + features = extract_function_features( + viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x401563) + ) + assert capa.features.Characteristic("cross section flow", True) 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', True) in features + features = extract_function_features( + viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA6FEC) + ) + assert capa.features.Characteristic("fs access", True) 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 + 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', True) in features + assert capa.features.Characteristic("embedded pe", True) in features def test_stackstring_features(mimikatz): features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x4556E5)) - assert capa.features.Characteristic('stack string', True) in features + assert capa.features.Characteristic("stack string", True) in features def test_switch_features(mimikatz): features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x409411)) - assert capa.features.Characteristic('switch', True) in features + assert capa.features.Characteristic("switch", True) in features features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x409393)) - assert capa.features.Characteristic('switch', True) not in features + assert capa.features.Characteristic("switch", True) not 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', True) 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", True) in features - features = extract_function_features(viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10007B00)) - assert capa.features.Characteristic('recursive call', True) not in features + features = extract_function_features( + viv_utils.Function( + sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, + 0x10007B00, + ) + ) + assert capa.features.Characteristic("recursive call", True) 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', True) in features +def test_loop_feature( + sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41, +): + features = extract_function_features( + viv_utils.Function( + sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, + 0x10003D30, + ) + ) + assert capa.features.Characteristic("loop", True) in features - features = extract_function_features(viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10007250)) - assert capa.features.Characteristic('loop', True) not in features + features = extract_function_features( + viv_utils.Function( + sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, + 0x10007250, + ) + ) + assert capa.features.Characteristic("loop", True) 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 + 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', True) in features - assert len(features[capa.features.Characteristic('calls to', True)]) == 1 + features = extract_function_features( + viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60) + ) + assert capa.features.Characteristic("calls to", True) in features + assert len(features[capa.features.Characteristic("calls to", True)]) == 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', True) in features - assert len(features[capa.features.Characteristic('calls to', True)]) == 8 + features = extract_function_features( + viv_utils.Function(sample_lab21_01.vw, 0x1400052D0) + ) # memcpy + assert capa.features.Characteristic("calls to", True) in features + assert len(features[capa.features.Characteristic("calls to", True)]) == 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', True) in features - assert len(features[capa.features.Characteristic('calls from', True)]) == 23 + features = extract_function_features( + viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60) + ) + assert capa.features.Characteristic("calls from", True) in features + assert len(features[capa.features.Characteristic("calls from", True)]) == 23 def test_basic_block_count(sample_9324d1a8ae37a36ae560c37448c9705a): - features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60)) + 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', True) in features - assert len(features[capa.features.Characteristic('indirect call', True)]) == 3 + features = extract_function_features( + viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA68A0) + ) + assert capa.features.Characteristic("indirect call", True) in features + assert len(features[capa.features.Characteristic("indirect call", True)]) == 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 + 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