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

Code before:
            if not isinstance(feature.value, str):
                # this is a programming error: String should only contain str
                raise ValueError("unexpected feature value type")

            if self.value in feature.value:
                matches[feature.value].update(locations)
                if short_circuit:
                    # we found one matching string, thats sufficient to match.
                    # don't collect other matching strings in this mode.
                    break

        if matches:
            # collect all locations
            locations = set()
            for locs in matches.values():
                locations.update(locs)

            # unlike other features, we cannot return put a reference to `self` directly in a `Result`.
            # this is because `self` may match on many strings, so we can't stuff the matched value into it.
            # instead, return a new instance that has a reference to both the substring and the matched values.
            return Result(True, _MatchedSubstring(self, dict(matches)), [], locations=locations)
        else:
            return Result(False, _MatchedSubstring(self, {}), [])

    def get_value_str(self) -> str:
        return escape_string(self.value)

    def __str__(self):
        assert isinstance(self.value, str)
        return "substring(%s)" % escape_string(self.value)


class _MatchedSubstring(Substring):
    """
    this represents specific match instances of a substring feature.
    treat it the same as a `Substring` except it has the `matches` field that contains the complete strings that matched.

    note: this type should only ever be constructed by `Substring.evaluate()`. it is not part of the public API.
    """

    def __init__(self, substring: Substring, matches: Dict[str, Set[Address]]):
        """
        args:
          substring: the substring feature that matches.
          match: mapping from matching string to its locations.
        """
        super().__init__(str(substring.value), description=substring.description)
        # we want this to collide with the name of `Substring` above,
        # so that it works nicely with the renderers.
        self.name = "substring"
        # this may be None if the substring doesn't match

Review comment (security):
mypy:
```
capa/features/common.py:183: error: Argument 1 to "escape_string" has incompatible type "Union[str, int, float, bytes]"; expected "str"  [arg-type]
capa/features/common.py:239: error: Argument 1 to "escape_string" has incompatible type "Union[str, int, float, bytes]"; expected "str"  [arg-type]
Found 2 errors in 1 file (checked 109 source files)
```

Code after:

            if not isinstance(feature.value, str):
                # this is a programming error: String should only contain str
                raise ValueError("unexpected feature value type")

            if self.value in feature.value:
                matches[feature.value].update(locations)
                if short_circuit:
                    # we found one matching string, thats sufficient to match.
                    # don't collect other matching strings in this mode.
                    break

        if matches:
            # collect all locations
            locations = set()
            for locs in matches.values():
                locations.update(locs)

            # unlike other features, we cannot return put a reference to `self` directly in a `Result`.
            # this is because `self` may match on many strings, so we can't stuff the matched value into it.
            # instead, return a new instance that has a reference to both the substring and the matched values.
            return Result(True, _MatchedSubstring(self, dict(matches)), [], locations=locations)
        else:
            return Result(False, _MatchedSubstring(self, {}), [])

    def get_value_str(self) -> str:
        assert isinstance(self.value, str)
        return escape_string(self.value)

    def __str__(self):
        assert isinstance(self.value, str)
        return "substring(%s)" % escape_string(self.value)


class _MatchedSubstring(Substring):
    """
    this represents specific match instances of a substring feature.
    treat it the same as a `Substring` except it has the `matches` field that contains the complete strings that matched.

    note: this type should only ever be constructed by `Substring.evaluate()`. it is not part of the public API.
    """

    def __init__(self, substring: Substring, matches: Dict[str, Set[Address]]):
        """
        args:
          substring: the substring feature that matches.
          match: mapping from matching string to its locations.
        """
        super().__init__(str(substring.value), description=substring.description)
        # we want this to collide with the name of `Substring` above,
        # so that it works nicely with the renderers.