From 7cb4ea927334a071fbe95c27f00fe95bd540f043 Mon Sep 17 00:00:00 2001 From: Pratham Chauhan Date: Thu, 30 Mar 2023 10:35:31 +0530 Subject: [PATCH] Fix lint issues --- capa/features/extractors/common.py | 22 ++++++++++++++++++---- capa/main.py | 2 +- capa/render/result_document.py | 19 ++++++++----------- tests/test_result_document.py | 8 +++++--- 4 files changed, 32 insertions(+), 19 deletions(-) diff --git a/capa/features/extractors/common.py b/capa/features/extractors/common.py index 806f895c..8a2c7899 100644 --- a/capa/features/extractors/common.py +++ b/capa/features/extractors/common.py @@ -10,16 +10,30 @@ import capa.features import capa.features.extractors.elf import capa.features.extractors.pefile import capa.features.extractors.strings -from capa.features.common import OS, FORMAT_PE, FORMAT_ELF, OS_ANY, OS_WINDOWS, FORMAT_FREEZE, FORMAT_RESULT, ARCH_ANY, Arch, Format, String, Feature +from capa.features.common import ( + OS, + OS_ANY, + ARCH_ANY, + FORMAT_PE, + FORMAT_ELF, + OS_WINDOWS, + FORMAT_FREEZE, + FORMAT_RESULT, + Arch, + Format, + String, + Feature, +) from capa.features.freeze import is_freeze from capa.features.address import NO_ADDRESS, Address, FileOffsetAddress logger = logging.getLogger(__name__) -#match strings for formats +# match strings for formats MATCH_PE = b"MZ" MATCH_ELF = b"\x7fELF" -MATCH_RESULT = b"{\"meta\":" +MATCH_RESULT = b'{"meta":' + def extract_file_strings(buf, **kwargs) -> Iterator[Tuple[String, Address]]: """ @@ -57,7 +71,7 @@ def extract_arch(buf) -> Iterator[Tuple[Feature, Address]]: elif buf.startswith(MATCH_ELF): with contextlib.closing(io.BytesIO(buf)) as f: arch = capa.features.extractors.elf.detect_elf_arch(f) - + elif buf.startswith(MATCH_RESULT): arch = ARCH_ANY diff --git a/capa/main.py b/capa/main.py index 081bd1fe..eb3c2bee 100644 --- a/capa/main.py +++ b/capa/main.py @@ -65,7 +65,7 @@ from capa.features.common import ( FORMAT_SC64, FORMAT_DOTNET, FORMAT_FREEZE, - FORMAT_RESULT + FORMAT_RESULT, ) from capa.features.address import NO_ADDRESS, Address from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor diff --git a/capa/render/result_document.py b/capa/render/result_document.py index b883dada..eac58b0b 100644 --- a/capa/render/result_document.py +++ b/capa/render/result_document.py @@ -5,8 +5,8 @@ # 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 datetime import json +import datetime from typing import Any, Dict, List, Tuple, Union, Optional from pydantic import Field, BaseModel @@ -152,18 +152,14 @@ class Metadata(FrozenModel): }, "feature_counts": { "file": self.analysis.feature_counts.file, - "functions": { - fc.address.to_capa(): fc.count for fc in self.analysis.feature_counts.functions - }, - }, - "library_functions": { - lf.address.to_capa(): lf.name for lf in self.analysis.library_functions + "functions": {fc.address.to_capa(): fc.count for fc in self.analysis.feature_counts.functions}, }, + "library_functions": {lf.address.to_capa(): lf.name for lf in self.analysis.library_functions}, }, } return capa_meta - + class CompoundStatementType: AND = "and" @@ -580,7 +576,7 @@ class ResultDocument(BaseModel): ) return ResultDocument(meta=Metadata.from_capa(meta), rules=rule_matches) - + def to_capa(self) -> Tuple[Dict, Dict]: meta = self.meta.to_capa() capabilities: Dict[str, List[Tuple[frz.Address, capa.features.common.Result]]] = {} @@ -607,10 +603,11 @@ class ResultDocument(BaseModel): statement=statement, success=match.success, locations=[frz.Address.to_capa(loc) for loc in match.locations], - children=[]) + children=[], + ) if rule_name not in capabilities: capabilities[rule_name] = [] capabilities[rule_name].append((frz.Address.from_capa(addr), result)) - return meta, capabilities \ No newline at end of file + return meta, capabilities diff --git a/tests/test_result_document.py b/tests/test_result_document.py index ccd79d72..da45d092 100644 --- a/tests/test_result_document.py +++ b/tests/test_result_document.py @@ -230,11 +230,13 @@ def test_basic_block_node_from_capa(): assert isinstance(node, rdoc.FeatureNode) assert isinstance(node.feature, frzf.BasicBlockFeature) + def test_json_to_rdoc(): path = fixtures.get_data_path_by_name("pma01-01-rd") - assert isinstance(rdoc.ResultDocument.parse_file(path),rdoc.ResultDocument) + assert isinstance(rdoc.ResultDocument.parse_file(path), rdoc.ResultDocument) + def test_rdoc_to_capa(): path = fixtures.get_data_path_by_name("pma01-01-rd") - assert len(rdoc.ResultDocument.parse_file(path).to_capa()) ==2 - assert isinstance(rdoc.ResultDocument.parse_file(path).to_capa(),tuple) \ No newline at end of file + assert len(rdoc.ResultDocument.parse_file(path).to_capa()) == 2 + assert isinstance(rdoc.ResultDocument.parse_file(path).to_capa(), tuple)