Repository: mandiant/capa
File: capa/features/extractors/base_extractor.py (Python)

Code before:
        """
        extract instruction-scope features.
        the arguments are opaque values previously provided by `.get_functions()`, etc.

        example::

            extractor = VivisectFeatureExtractor(vw, path)
            for function in extractor.get_functions():
                for bb in extractor.get_basic_blocks(function):
                    for insn in extractor.get_instructions(function, bb):
                        for feature, address in extractor.extract_insn_features(function, bb, insn):
                            print('0x%x: %s', address, feature)

        args:
          f [FunctionHandle]: an opaque value previously fetched from `.get_functions()`.
          bb [BBHandle]: an opaque value previously fetched from `.get_basic_blocks()`.
          insn [InsnHandle]: an opaque value previously fetched from `.get_instructions()`.

        yields:
          Tuple[Feature, Address]: feature and its location
        """
        raise NotImplementedError()


def FunctionFilter(extractor: StaticFeatureExtractor, functions: Set) -> StaticFeatureExtractor:
    get_functions = extractor.get_functions  # fetch original get_functions()

    def filtered_get_functions(self):
        yield from (f for f in get_functions() if f.address in functions)

    # make a copy of the extractor before decorating the get_functions() method
    new_extractor = copy(extractor)
    new_extractor.get_functions = MethodType(filtered_get_functions, extractor)  # type: ignore

    return new_extractor


@dataclass
class ProcessHandle:
    """
    reference to a process extracted by the sandbox.

    Attributes:
        address: process's address (pid)
        inner: sandbox-specific data
    """

    address: ProcessAddress
    inner: Any



Review comment (refactor):
```suggestion
    original_get_functions = extractor.get_functions  # fetch original get_functions()
```

Code after:
        """
        extract instruction-scope features.
        the arguments are opaque values previously provided by `.get_functions()`, etc.

        example::

            extractor = VivisectFeatureExtractor(vw, path)
            for function in extractor.get_functions():
                for bb in extractor.get_basic_blocks(function):
                    for insn in extractor.get_instructions(function, bb):
                        for feature, address in extractor.extract_insn_features(function, bb, insn):
                            print('0x%x: %s', address, feature)

        args:
          f [FunctionHandle]: an opaque value previously fetched from `.get_functions()`.
          bb [BBHandle]: an opaque value previously fetched from `.get_basic_blocks()`.
          insn [InsnHandle]: an opaque value previously fetched from `.get_instructions()`.

        yields:
          Tuple[Feature, Address]: feature and its location
        """
        raise NotImplementedError()


def FunctionFilter(extractor: StaticFeatureExtractor, functions: Set) -> StaticFeatureExtractor:
    original_get_functions = extractor.get_functions

    def filtered_get_functions(self):
        yield from (f for f in original_get_functions() if f.address in functions)

    # we make a copy of the original extractor object and then update its get_functions() method with the decorated filter one.
    # this is in order to preserve the original extractor object's get_functions() method, in case it is used elsewhere in the code.
    # an example where this is important is in our testfiles where we may use the same extractor object with different tests,
    # with some of these tests needing to install a functions filter on the extractor object.
    new_extractor = copy(extractor)
    new_extractor.get_functions = MethodType(filtered_get_functions, extractor)  # type: ignore

    return new_extractor


@dataclass
class ProcessHandle:
    """
    reference to a process extracted by the sandbox.

    Attributes:
        address: process's address (pid)
        inner: sandbox-specific data
    """

    address: ProcessAddress