Merge pull request #1396 from ooprathamm/read-render

Towards improving read and rendering of results

CHANGELOG.md

@@ -4,10 +4,10 @@
 
 ### New Features
 - add protobuf format for result documents #1219 @williballenthin @mr-tz 
-
 - extractor: add Binary Ninja feature extractor @xusheng6
 - new cli flag `--os` to override auto-detected operating system for a sample @captainGeech42
 - Change colour/highlight to "cyan" instead of "blue" for easy noticing.#1384 @ggold7046
+- add new format to parse output json back to capa #1396 @ooprathamm
 
 ### Breaking Changes
 

capa/features/common.py

@@ -450,6 +450,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

@@ -12,11 +12,14 @@
 import capa.features.extractors.strings
 from capa.features.common import (
  OS,
+ OS_ANY,
  OS_AUTO,
+ ARCH_ANY,
  FORMAT_PE,
  FORMAT_ELF,
  OS_WINDOWS,
  FORMAT_FREEZE,
+ FORMAT_RESULT,
  Arch,
  Format,
  String,
@@ -27,6 +30,11 @@
 
 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]]:
  """
@@ -40,12 +48,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)
@@ -56,10 +66,13 @@ 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(MATCH_RESULT):
+ yield Arch(ARCH_ANY), NO_ADDRESS
 
- 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)
 
@@ -88,9 +101,11 @@ def extract_os(buf, os=OS_AUTO) -> Iterator[Tuple[Feature, Address]]:
  if os != OS_AUTO:
  yield OS(os), NO_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

@@ -69,6 +69,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
@@ -1180,8 +1181,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_file(args.sample)
+ meta, capabilities = result_doc.to_capa()
+ elif format_ == FORMAT_FREEZE:
  with open(args.sample, "rb") as f:
  extractor = capa.features.freeze.load(f.read())
  else:
@@ -1217,17 +1220,18 @@ def main(argv=None):
  log_unsupported_os_error()
  return E_INVALID_FILE_OS
 
- meta = collect_metadata(argv, args.sample, args.format, args.os, args.rules, extractor)
+ if format_ != FORMAT_RESULT:
+ meta = collect_metadata(argv, args.sample, args.format, args.os, 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,7 @@
 # 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
+from typing import Any, Dict, List, Tuple, Union, Optional
 
 from pydantic import Field, BaseModel
 
@@ -125,6 +125,41 @@ 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"
@@ -543,3 +578,38 @@ def from_capa(cls, meta, rules: RuleSet, capabilities: MatchResults) -> "ResultD
  )
 
  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]]] = {}
+
+ for rule_name, rule_match in self.rules.items():
+ # Parse the YAML source into a Rule instance
+ rule = capa.rules.Rule.from_yaml(rule_match.source)
+
+ # 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 = rule.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/fixtures.py

@@ -241,6 +241,8 @@ def get_data_path_by_name(name):
  return os.path.join(CD, "data", "kernel32-64.dll_")
  elif name == "pma01-01":
  return os.path.join(CD, "data", "Practical Malware Analysis Lab 01-01.dll_")
+ elif name == "pma01-01-rd":
+ return os.path.join(CD, "data", "rd", "Practical Malware Analysis Lab 01-01.dll_.json")
  elif name == "pma12-04":
  return os.path.join(CD, "data", "Practical Malware Analysis Lab 12-04.exe_")
  elif name == "pma16-01":

tests/test_result_document.py

@@ -8,6 +8,7 @@
 import copy
 
 import pytest
+import fixtures
 from fixtures import *
 
 import capa
@@ -268,3 +269,14 @@ def assert_round_trip(rd: rdoc.ResultDocument):
 def test_round_trip(request, rd_file):
  rd: rdoc.ResultDocument = request.getfixturevalue(rd_file)
  assert_round_trip(rd)
+
+
+def test_json_to_rdoc():
+ path = fixtures.get_data_path_by_name("pma01-01-rd")
+ 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)