fix mypy issues and bugs

This commit is contained in:
Yacine Elhamer
2023-07-10 14:11:10 +01:00
parent 37e4b913b0
commit af256bc0e9
3 changed files with 8 additions and 8 deletions

View File

@@ -74,7 +74,7 @@ class ThreadAddress(Address):
return f"thread(tid: {self.tid})"
def __hash__(self):
return hash((self.ppid, self.pid, self.tid))
return hash((self.process, self.tid))
def __eq__(self, other):
assert isinstance(other, ThreadAddress)

View File

@@ -2,7 +2,7 @@ from typing import Dict, List, Tuple, Union, TypeAlias
from dataclasses import dataclass
from capa.features.common import Feature
from capa.features.address import NO_ADDRESS, Address
from capa.features.address import NO_ADDRESS, Address, ThreadAddress, ProcessAddress
from capa.features.extractors.base_extractor import (
BBHandle,
InsnHandle,
@@ -88,7 +88,7 @@ class ThreadFeatures:
@dataclass
class ProcessFeatures:
features: List[Tuple[Address, Feature]]
threads: Dict[Address, ThreadFeatures]
threads: Dict[ThreadAddress, ThreadFeatures]
@dataclass
@@ -96,7 +96,7 @@ class NullDynamicFeatureExtractor(DynamicFeatureExtractor):
base_address: Address
global_features: List[Feature]
file_features: List[Tuple[Address, Feature]]
processes: Dict[Address, ProcessFeatures]
processes: Dict[ProcessAddress, ProcessFeatures]
def extract_global_features(self):
for feature in self.global_features:
@@ -108,7 +108,7 @@ class NullDynamicFeatureExtractor(DynamicFeatureExtractor):
def get_processes(self):
for address in sorted(self.processes.keys()):
yield ProcessHandle(address=address, inner={}, pid=address.pid)
yield ProcessHandle(address=address, inner={})
def extract_process_features(self, p):
for addr, feature in self.processes[p.address].features:
@@ -116,7 +116,7 @@ class NullDynamicFeatureExtractor(DynamicFeatureExtractor):
def get_threads(self, p):
for address in sorted(self.processes[p].threads.keys()):
yield ThreadHandle(address=address, inner={}, tid=address.pid)
yield ThreadHandle(address=address, inner={})
def extract_thread_features(self, p, t):
for addr, feature in self.processes[p.address].threads[t.address].features:

View File

@@ -50,7 +50,7 @@ class AddressType(str, Enum):
class Address(HashableModel):
type: AddressType
value: Union[int, Tuple[int, int], Tuple[int, int, int], None]
value: Union[int, Tuple[int, ...], None]
@classmethod
def from_capa(cls, a: capa.features.address.Address) -> "Address":
@@ -73,7 +73,7 @@ class Address(HashableModel):
return cls(type=AddressType.PROCESS, value=(a.ppid, a.pid))
elif isinstance(a, capa.features.address.ThreadAddress):
return cls(type=AddressType.THREAD, value=(a.ppid, a.pid, a.tid))
return cls(type=AddressType.THREAD, value=(a.process.ppid, a.process.pid, a.tid))
elif isinstance(a, capa.features.address.DynamicAddress):
return cls(type=AddressType.DYNAMIC, value=(a.id, a.return_address))