capa/tests/test_main.py

import textwrap

import capa.main
import capa.rules
import capa.engine
from capa.engine import *
import capa.features
import capa.features.extractors.viv

from fixtures import *


def test_main(sample_9324d1a8ae37a36ae560c37448c9705a):
    # tests rules can be loaded successfully
    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_file.write(RULE_CONTENT)
    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
    )


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
                    """
                )
            ),
        ]
    )
    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,
    )
    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


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,
    )
    capabilities = capa.main.find_capabilities(rules, extractor)
    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
                    """
                )
            )
        ]
    )
    # tight loop at 0x403685
    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(
        sample_9324d1a8ae37a36ae560c37448c9705a.vw,
        sample_9324d1a8ae37a36ae560c37448c9705a.path,
    )
    capabilities = capa.main.find_capabilities(rules, extractor)
    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
                    """
                )
            )
        ]
    )

    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