diff --git a/capa/engine.py b/capa/engine.py index b5fbb412..dde1e7c8 100644 --- a/capa/engine.py +++ b/capa/engine.py @@ -8,7 +8,7 @@ import copy import collections -from typing import TYPE_CHECKING, Set, Dict, List, Tuple, Mapping, Iterable, Iterator, Union, cast +from typing import TYPE_CHECKING, Set, Dict, List, Tuple, Union, Mapping, Iterable, Iterator, cast import capa.perf import capa.features.common diff --git a/capa/features/extractors/dnfile_.py b/capa/features/extractors/dnfile_.py index cf82bbce..03859db7 100644 --- a/capa/features/extractors/dnfile_.py +++ b/capa/features/extractors/dnfile_.py @@ -109,7 +109,7 @@ class DnfileFeatureExtractor(FeatureExtractor): assert self.pe is not None assert self.pe.net is not None assert self.pe.net.struct is not None - + return self.pe.net.struct.MajorRuntimeVersion, self.pe.net.struct.MinorRuntimeVersion def get_meta_version_string(self) -> str: diff --git a/capa/features/extractors/null.py b/capa/features/extractors/null.py index f8d6d077..892eadc8 100644 --- a/capa/features/extractors/null.py +++ b/capa/features/extractors/null.py @@ -68,10 +68,5 @@ class NullFeatureExtractor(FeatureExtractor): yield InsnHandle(address, None) def extract_insn_features(self, f, bb, insn): - for address, feature in ( - self.functions[f.address] - .basic_blocks[bb.address] - .instructions[insn.address] - .features - ): + for address, feature in self.functions[f.address].basic_blocks[bb.address].instructions[insn.address].features: yield feature, address diff --git a/capa/ida/plugin/form.py b/capa/ida/plugin/form.py index e6deecb7..eb205ae6 100644 --- a/capa/ida/plugin/form.py +++ b/capa/ida/plugin/form.py @@ -11,7 +11,7 @@ import copy import logging import itertools import collections -from typing import Set, Dict, Optional, List, Any +from typing import Any, Set, Dict, List, Optional import idaapi import ida_kernwin @@ -1108,7 +1108,7 @@ class CapaExplorerForm(idaapi.PluginForm): _, file_matches = capa.engine.match( capa.rules.RuleSet(list(capa.rules.get_rules_and_dependencies(rules, rule.name))).file_rules, file_features, - NO_ADDRESS + NO_ADDRESS, ) except Exception as e: self.set_rulegen_status("Failed to match rule (%s)" % e) diff --git a/capa/ida/plugin/view.py b/capa/ida/plugin/view.py index 75abf59c..0f577c7d 100644 --- a/capa/ida/plugin/view.py +++ b/capa/ida/plugin/view.py @@ -18,7 +18,7 @@ import capa.ida.helpers import capa.features.common import capa.features.basicblock from capa.ida.plugin.item import CapaExplorerFunctionItem -from capa.features.address import _NoAddress, AbsoluteVirtualAddress +from capa.features.address import AbsoluteVirtualAddress, _NoAddress from capa.ida.plugin.model import CapaExplorerDataModel MAX_SECTION_SIZE = 750 diff --git a/tests/test_engine.py b/tests/test_engine.py index 8fee9b92..89c3b739 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -5,16 +5,17 @@ # Unless required by applicable law or agreed to in writing, software distributed under the License # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and limitations under the License. +import capa.features.address from capa.engine import * from capa.features import * from capa.features.insn import * -import capa.features.address ADDR1 = capa.features.address.AbsoluteVirtualAddress(0x401001) ADDR2 = capa.features.address.AbsoluteVirtualAddress(0x401002) ADDR3 = capa.features.address.AbsoluteVirtualAddress(0x401003) ADDR4 = capa.features.address.AbsoluteVirtualAddress(0x401004) + def test_number(): assert Number(1).evaluate({Number(0): {ADDR1}}) == False assert Number(1).evaluate({Number(1): {ADDR1}}) == True @@ -50,7 +51,12 @@ def test_some(): assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {ADDR1}}) == False assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {ADDR1}, Number(1): {ADDR1}}) == False - assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {ADDR1}, Number(1): {ADDR1}, Number(2): {ADDR1}}) == True + assert ( + Some(2, [Number(1), Number(2), Number(3)]).evaluate( + {Number(0): {ADDR1}, Number(1): {ADDR1}, Number(2): {ADDR1}} + ) + == True + ) assert ( Some(2, [Number(1), Number(2), Number(3)]).evaluate( {Number(0): {ADDR1}, Number(1): {ADDR1}, Number(2): {ADDR1}, Number(3): {ADDR1}} diff --git a/tests/test_rules.py b/tests/test_rules.py index d5aea406..43d7002d 100644 --- a/tests/test_rules.py +++ b/tests/test_rules.py @@ -13,10 +13,10 @@ import pytest import capa.rules import capa.engine import capa.features.common -from capa.features.address import AbsoluteVirtualAddress +import capa.features.address +from capa.engine import Or from capa.features.file import FunctionName from capa.features.insn import Number, Offset, Property -from capa.engine import Or from capa.features.common import ( OS, OS_LINUX, @@ -31,8 +31,7 @@ from capa.features.common import ( Substring, FeatureAccess, ) -import capa.features.address - +from capa.features.address import AbsoluteVirtualAddress ADDR1 = capa.features.address.AbsoluteVirtualAddress(0x401001) ADDR2 = capa.features.address.AbsoluteVirtualAddress(0x401002)