diff --git a/capa/features/common.py b/capa/features/common.py index b32caa7c..899d3143 100644 --- a/capa/features/common.py +++ b/capa/features/common.py @@ -10,6 +10,7 @@ import re import codecs import logging import collections +from typing import Dict, List, Union import capa.engine import capa.features.common @@ -27,16 +28,16 @@ ARCH_X64 = "x64" VALID_ARCH = (ARCH_X32, ARCH_X64) -def bytes_to_str(b): +def bytes_to_str(b: bytes) -> str: return str(codecs.encode(b, "hex").decode("utf-8")) -def hex_string(h): +def hex_string(h: str) -> str: """render hex string e.g. "0a40b1" as "0A 40 B1" """ return " ".join(h[i : i + 2] for i in range(0, len(h), 2)).upper() -def escape_string(s): +def escape_string(s: str) -> str: """escape special characters""" s = repr(s) if not s.startswith(('"', "'")): @@ -51,7 +52,7 @@ def escape_string(s): class Feature(object): - def __init__(self, value, arch=None, description=None): + def __init__(self, value: Union[str, int, bytes], arch=None, description=None): """ Args: value (any): the value of the feature, such as the number or string. @@ -79,14 +80,14 @@ class Feature(object): def __eq__(self, other): return self.name == other.name and self.value == other.value and self.arch == other.arch - def get_value_str(self): + def get_value_str(self) -> str: """ render the value of this feature, for use by `__str__` and friends. subclasses should override to customize the rendering. Returns: any """ - return self.value + return str(self.value) def __str__(self): if self.value is not None: @@ -100,7 +101,7 @@ class Feature(object): def __repr__(self): return str(self) - def evaluate(self, ctx): + def evaluate(self, ctx: Dict["Feature", List[int]]) -> capa.engine.Result: return capa.engine.Result(self in ctx, self, [], locations=ctx.get(self, [])) def freeze_serialize(self): @@ -123,24 +124,26 @@ class Feature(object): class MatchedRule(Feature): - def __init__(self, value, description=None): + def __init__(self, value: str, description=None): super(MatchedRule, self).__init__(value, description=description) self.name = "match" class Characteristic(Feature): - def __init__(self, value, description=None): + def __init__(self, value: str, description=None): super(Characteristic, self).__init__(value, description=description) class String(Feature): - def __init__(self, value, description=None): + def __init__(self, value: str, description=None): super(String, self).__init__(value, description=description) class Regex(String): - def __init__(self, value, description=None): + def __init__(self, value: str, description=None): super(Regex, self).__init__(value, description=description) + self.value = value + pat = self.value[len("/") : -len("/")] flags = re.DOTALL if value.endswith("/i"): @@ -202,7 +205,7 @@ class _MatchedRegex(Regex): note: this type should only ever be constructed by `Regex.evaluate()`. it is not part of the public API. """ - def __init__(self, regex, matches): + def __init__(self, regex: Regex, matches): """ args: regex (Regex): the regex feature that matches. @@ -223,15 +226,16 @@ class _MatchedRegex(Regex): class StringFactory(object): - def __new__(cls, value, description=None): + def __new__(cls, value: str, description=None): if value.startswith("/") and (value.endswith("/") or value.endswith("/i")): return Regex(value, description=description) return String(value, description=description) class Bytes(Feature): - def __init__(self, value, description=None): + def __init__(self, value: bytes, description=None): super(Bytes, self).__init__(value, description=description) + self.value = value def evaluate(self, ctx): for feature, locations in ctx.items(): diff --git a/capa/features/extractors/viv/indirect_calls.py b/capa/features/extractors/viv/indirect_calls.py index 6a655596..a090355a 100644 --- a/capa/features/extractors/viv/indirect_calls.py +++ b/capa/features/extractors/viv/indirect_calls.py @@ -7,7 +7,7 @@ # See the License for the specific language governing permissions and limitations under the License. import collections -from typing import List, Tuple, Optional +from typing import TYPE_CHECKING, List, Tuple, Optional import envi import vivisect.const @@ -15,7 +15,6 @@ import envi.archs.i386.disasm import envi.archs.amd64.disasm from vivisect import VivWorkspace -from typing import TYPE_CHECKING if TYPE_CHECKING: from capa.features.extractors.viv.extractor import InstructionHandle diff --git a/capa/rules.py b/capa/rules.py index 2d29e96b..fb2eb877 100644 --- a/capa/rules.py +++ b/capa/rules.py @@ -20,9 +20,10 @@ try: except ImportError: from backports.functools_lru_cache import lru_cache +from typing import Any, Set, Dict, List, Union, Iterator + import yaml import ruamel.yaml -from typing import List, Any, Union, Dict, Iterator, Set import capa.rules import capa.engine @@ -32,8 +33,8 @@ import capa.features.file import capa.features.insn import capa.features.common import capa.features.basicblock -from capa.features.common import MAX_BYTES_FEATURE_SIZE, Feature from capa.rules import Rule +from capa.features.common import MAX_BYTES_FEATURE_SIZE, Feature logger = logging.getLogger(__name__)