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

Code before:
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_call_features(
        self, ph: ProcessHandle, th: ThreadHandle, ch: CallHandle
    ) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all features of a call. These include:
        - api name
        - bytes/strings/numbers extracted from arguments
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_call_name(self, ph: ProcessHandle, th: ThreadHandle, ch: CallHandle) -> str:
        """
        Returns the human-readable name for the given call,
        such as as rendered API log entry, like:

            Foo(1, "two", b"\x00\x11") -> -1
        """
        raise NotImplementedError()


def ProcessFilter(extractor: DynamicFeatureExtractor, processes: Set) -> DynamicFeatureExtractor:
    original_get_processes = extractor.get_processes  # fetch original get_processes()

    def filtered_get_processes(self):
        yield from (f for f in original_get_processes() if f.address.pid in processes)

    # we make a copy of the original extractor object and then update its get_processes() method with the decorated filter one.
    # this is in order to preserve the original extractor object's get_processes() 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 processes filter on the extractor object.
    new_extractor = copy(extractor)
    new_extractor.get_processes = MethodType(filtered_get_processes, extractor)  # type: ignore

    return new_extractor


FeatureExtractor: TypeAlias = Union[StaticFeatureExtractor, DynamicFeatureExtractor]


Review comment (refactor):
```suggestion
    original_get_processes = extractor.get_processes
```

Code after:
        raise NotImplementedError()

    @abc.abstractmethod
    def extract_call_features(
        self, ph: ProcessHandle, th: ThreadHandle, ch: CallHandle
    ) -> Iterator[Tuple[Feature, Address]]:
        """
        Yields all features of a call. These include:
        - api name
        - bytes/strings/numbers extracted from arguments
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def get_call_name(self, ph: ProcessHandle, th: ThreadHandle, ch: CallHandle) -> str:
        """
        Returns the human-readable name for the given call,
        such as as rendered API log entry, like:

            Foo(1, "two", b"\x00\x11") -> -1
        """
        raise NotImplementedError()


def ProcessFilter(extractor: DynamicFeatureExtractor, processes: Set) -> DynamicFeatureExtractor:
    original_get_processes = extractor.get_processes

    def filtered_get_processes(self):
        yield from (f for f in original_get_processes() if f.address.pid in processes)

    # we make a copy of the original extractor object and then update its get_processes() method with the decorated filter one.
    # this is in order to preserve the original extractor object's get_processes() 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 processes filter on the extractor object.
    new_extractor = copy(extractor)
    new_extractor.get_processes = MethodType(filtered_get_processes, extractor)  # type: ignore

    return new_extractor


FeatureExtractor: TypeAlias = Union[StaticFeatureExtractor, DynamicFeatureExtractor]