Towards improving read and rendering of results

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## master (unreleased)
 
 ### New Features
+- add new format to parse output json back to capa #1396 @ooprathamm
 
 ### Breaking Changes
 

capa/features/common.py

@@ -448,6 +448,7 @@ def evaluate(self, ctx, **kwargs):
 FORMAT_SC32 = "sc32"
 FORMAT_SC64 = "sc64"
 FORMAT_FREEZE = "freeze"
+FORMAT_RESULT = "result"
 FORMAT_UNKNOWN = "unknown"
 
 

capa/features/extractors/common.py

@@ -10,12 +10,16 @@
 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_WINDOWS, FORMAT_FREEZE, Arch, Format, String, Feature
+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.freeze import is_freeze
 from capa.features.address import NO_ADDRESS, Address, FileOffsetAddress
 
 logger = logging.getLogger(__name__)
 
+#match strings for formats
+MATCH_PE = b"MZ"
+MATCH_ELF = b"\x7fELF"
+MATCH_RESULT = b"{\"meta\":"
 
 def extract_file_strings(buf, **kwargs) -> Iterator[Tuple[String, Address]]:
     """
@@ -29,12 +33,14 @@ def extract_file_strings(buf, **kwargs) -> Iterator[Tuple[String, Address]]:
 
 
 def extract_format(buf) -> Iterator[Tuple[Feature, Address]]:
-    if buf.startswith(b"MZ"):
+    if buf.startswith(MATCH_PE):
         yield Format(FORMAT_PE), NO_ADDRESS
-    elif buf.startswith(b"\x7fELF"):
+    elif buf.startswith(MATCH_ELF):
         yield Format(FORMAT_ELF), NO_ADDRESS
     elif is_freeze(buf):
         yield Format(FORMAT_FREEZE), NO_ADDRESS
+    elif buf.startswith(MATCH_RESULT):
+        yield Format(FORMAT_RESULT), NO_ADDRESS
     else:
         # we likely end up here:
         #  1. handling a file format (e.g. macho)
@@ -45,12 +51,15 @@ def extract_format(buf) -> Iterator[Tuple[Feature, Address]]:
 
 
 def extract_arch(buf) -> Iterator[Tuple[Feature, Address]]:
-    if buf.startswith(b"MZ"):
+    if buf.startswith(MATCH_PE):
         yield from capa.features.extractors.pefile.extract_file_arch(pe=pefile.PE(data=buf))
 
-    elif buf.startswith(b"\x7fELF"):
+    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
 
         if arch not in capa.features.common.VALID_ARCH:
             logger.debug("unsupported arch: %s", arch)
@@ -74,9 +83,11 @@ def extract_arch(buf) -> Iterator[Tuple[Feature, Address]]:
 
 
 def extract_os(buf) -> Iterator[Tuple[Feature, Address]]:
-    if buf.startswith(b"MZ"):
+    if buf.startswith(MATCH_PE):
         yield OS(OS_WINDOWS), NO_ADDRESS
-    elif buf.startswith(b"\x7fELF"):
+    elif buf.startswith(MATCH_RESULT):
+        yield OS(OS_ANY), NO_ADDRESS
+    elif buf.startswith(MATCH_ELF):
         with contextlib.closing(io.BytesIO(buf)) as f:
             os = capa.features.extractors.elf.detect_elf_os(f)
 

capa/main.py

@@ -65,6 +65,7 @@
     FORMAT_SC64,
     FORMAT_DOTNET,
     FORMAT_FREEZE,
+    FORMAT_RESULT
 )
 from capa.features.address import NO_ADDRESS, Address
 from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor
@@ -1123,8 +1124,10 @@ def main(argv=None):
             if not (args.verbose or args.vverbose or args.json):
                 logger.debug("file limitation short circuit, won't analyze fully.")
                 return E_FILE_LIMITATION
-
-    if format_ == FORMAT_FREEZE:
+    if format_ == FORMAT_RESULT:
+        result_doc = capa.render.result_document.ResultDocument.parse_raw(args.sample)
+        meta, capabilities = result_doc.to_capa(rules)
+    elif format_ == FORMAT_FREEZE:
         with open(args.sample, "rb") as f:
             extractor = capa.features.freeze.load(f.read())
     else:
@@ -1154,17 +1157,17 @@ def main(argv=None):
             log_unsupported_os_error()
             return E_INVALID_FILE_OS
 
-    meta = collect_metadata(argv, args.sample, args.rules, extractor)
+        meta = collect_metadata(argv, args.sample, args.rules, extractor)
 
-    capabilities, counts = find_capabilities(rules, extractor, disable_progress=args.quiet)
-    meta["analysis"].update(counts)
-    meta["analysis"]["layout"] = compute_layout(rules, extractor, capabilities)
+        capabilities, counts = find_capabilities(rules, extractor, disable_progress=args.quiet)
+        meta["analysis"].update(counts)
+        meta["analysis"]["layout"] = compute_layout(rules, extractor, capabilities)
 
-    if has_file_limitation(rules, capabilities):
-        # bail if capa encountered file limitation e.g. a packed binary
-        # do show the output in verbose mode, though.
-        if not (args.verbose or args.vverbose or args.json):
-            return E_FILE_LIMITATION
+        if has_file_limitation(rules, capabilities):
+            # bail if capa encountered file limitation e.g. a packed binary
+            # do show the output in verbose mode, though.
+            if not (args.verbose or args.vverbose or args.json):
+                return E_FILE_LIMITATION
 
     if args.json:
         print(capa.render.json.render(meta, rules, capabilities))

capa/render/result_document.py

@@ -6,7 +6,8 @@
 #  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
-from typing import Any, Dict, Tuple, Union, Optional
+import json
+from typing import Any, Dict, List, Tuple, Union, Optional
 
 from pydantic import Field, BaseModel
 
@@ -124,6 +125,45 @@ def from_capa(cls, meta: Any) -> "Metadata":
             ),
         )
 
+    def to_capa(self) -> Dict[str, Any]:
+        capa_meta = {
+            "timestamp": self.timestamp.isoformat(),
+            "version": self.version,
+            "sample": {
+                "md5": self.sample.md5,
+                "sha1": self.sample.sha1,
+                "sha256": self.sample.sha256,
+                "path": self.sample.path,
+            },
+            "analysis": {
+                "format": self.analysis.format,
+                "arch": self.analysis.arch,
+                "os": self.analysis.os,
+                "extractor": self.analysis.extractor,
+                "rules": self.analysis.rules,
+                "base_address": self.analysis.base_address.to_capa(),
+                "layout": {
+                    "functions": {
+                        f.address.to_capa(): {
+                            "matched_basic_blocks": [bb.address.to_capa() for bb in f.matched_basic_blocks]
+                        }
+                        for f in self.analysis.layout.functions
+                    }
+                },
+                "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
+                },
+            },
+        }
+
+        return capa_meta
+
 
 class CompoundStatementType:
     AND = "and"
@@ -540,3 +580,43 @@ def from_capa(cls, meta, rules: RuleSet, capabilities: MatchResults) -> "ResultD
             )
 
         return ResultDocument(meta=Metadata.from_capa(meta), rules=rule_matches)
+
+    @classmethod
+    def parse_raw(cls, path: str):
+        with open(path, "rb") as f:
+            buf= f.read()
+        data = json.loads(buf)
+        return ResultDocument(**data)    
+
+    def to_capa(self, rules: RuleSet) -> Tuple[Dict, Dict]:
+        meta = self.meta.to_capa()
+        capabilities: Dict[str, List[Tuple[frz.Address, capa.features.common.Result]]] ={}
+
+        for rule_name, rule_match in self.rules.items():
+
+            # Extract the capabilities from the RuleMatches object
+            for addr, match in rule_match.matches:
+
+                if isinstance(match.node, StatementNode):
+                    if isinstance(match.node.statement, CompoundStatement):
+                        statement = rules[rule_name].statement
+                    else:
+                        statement = statement_from_capa(match.node.statement)
+                elif isinstance(match.node, FeatureNode):
+                    statement = match.node.feature.to_capa()
+                    if isinstance(statement, (capa.features.common.String, capa.features.common.Regex)):
+                        statement.matches = match.captures
+                else:
+                    raise ValueError("Invalid node type")
+
+                result = capa.features.common.Result(
+                statement=statement,
+                success=match.success,
+                locations=[frz.Address.to_capa(loc) for loc in match.locations],
+                children=[])
+
+                if rule_name not in capabilities:
+                    capabilities[rule_name]=[]
+                capabilities[rule_name].append((frz.Address.from_capa(addr),result))
+
+        return meta, capabilities

tests/test_result_document.py

@@ -6,6 +6,8 @@
 #  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 fixtures
+
 import capa
 import capa.engine as ceng
 import capa.render.result_document as rdoc
@@ -227,3 +229,10 @@ def test_basic_block_node_from_capa():
     node = rdoc.node_from_capa(capa.features.basicblock.BasicBlock(""))
     assert isinstance(node, rdoc.FeatureNode)
     assert isinstance(node.feature, frzf.BasicBlockFeature)
+
+def test_json_to_rdoc (capsys, tmpdir):
+    path = fixtures.get_data_path_by_name("pma01-01")
+    assert capa.main.main([path, "-j"]) == 0
+    temp_file = tmpdir.join("capa.json")
+    temp_file.write(capsys.readouterr().out)
+    assert isinstance(rdoc.ResultDocument.parse_raw(temp_file),rdoc.ResultDocument)