From 06f8943bc41952fa4f874ced33830ea33894e672 Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Wed, 11 Aug 2021 09:10:04 -0600
Subject: [PATCH] features: add format/pe and format/elf characteristics

---
 capa/features/common.py |  7 +++++++
 capa/rules.py           |  8 ++++++++
 tests/test_rules.py     | 19 ++++++++++++++++++-
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/capa/features/common.py b/capa/features/common.py
index 7a995f81..5f6e19cb 100644
--- a/capa/features/common.py
+++ b/capa/features/common.py
@@ -32,6 +32,10 @@ OS_LINUX = "os/linux"
 OS_MACOS = "os/macos"
 VALID_OS = (OS_WINDOWS, OS_LINUX, OS_MACOS)
 
+FORMAT_PE = "format/pe"
+FORMAT_ELF = "format/elf"
+VALID_FORMAT = (FORMAT_PE, FORMAT_ELF)
+
 
 def bytes_to_str(b: bytes) -> str:
     return str(codecs.encode(b, "hex").decode("utf-8"))
@@ -144,6 +148,9 @@ CHARACTERISTIC_WINDOWS = Characteristic(OS_WINDOWS)
 CHARACTERISTIC_LINUX = Characteristic(OS_LINUX)
 CHARACTERISTIC_MACOS = Characteristic(OS_MACOS)
 
+CHARACTERISTIC_PE = Characteristic(FORMAT_PE)
+CHARACTERISTIC_ELF = Characteristic(FORMAT_ELF)
+
 
 class String(Feature):
     def __init__(self, value: str, description=None):
diff --git a/capa/rules.py b/capa/rules.py
index 6d4f0872..e3264998 100644
--- a/capa/rules.py
+++ b/capa/rules.py
@@ -36,6 +36,8 @@ import capa.features.basicblock
 from capa.engine import Statement, FeatureSet
 from capa.features.common import MAX_BYTES_FEATURE_SIZE, Feature
 from capa.features.common import CHARACTERISTIC_WINDOWS, CHARACTERISTIC_LINUX, CHARACTERISTIC_MACOS
+from capa.features.common import CHARACTERISTIC_PE, CHARACTERISTIC_ELF
+
 
 logger = logging.getLogger(__name__)
 
@@ -82,6 +84,8 @@ SUPPORTED_FEATURES = {
         CHARACTERISTIC_WINDOWS,
         CHARACTERISTIC_LINUX,
         CHARACTERISTIC_MACOS,
+        CHARACTERISTIC_PE,
+        CHARACTERISTIC_ELF,
     },
     FUNCTION_SCOPE: {
         # plus basic block scope features, see below
@@ -93,6 +97,8 @@ SUPPORTED_FEATURES = {
         CHARACTERISTIC_WINDOWS,
         CHARACTERISTIC_LINUX,
         CHARACTERISTIC_MACOS,
+        CHARACTERISTIC_PE,
+        CHARACTERISTIC_ELF,
     },
     BASIC_BLOCK_SCOPE: {
         capa.features.common.MatchedRule,
@@ -113,6 +119,8 @@ SUPPORTED_FEATURES = {
         CHARACTERISTIC_WINDOWS,
         CHARACTERISTIC_LINUX,
         CHARACTERISTIC_MACOS,
+        CHARACTERISTIC_PE,
+        CHARACTERISTIC_ELF,
     },
 }
 
diff --git a/tests/test_rules.py b/tests/test_rules.py
index 6b57b89a..abc37968 100644
--- a/tests/test_rules.py
+++ b/tests/test_rules.py
@@ -15,7 +15,7 @@ import capa.engine
 import capa.features.common
 from capa.features.file import FunctionName
 from capa.features.insn import Number, Offset
-from capa.features.common import ARCH_X32, ARCH_X64, OS_WINDOWS, String, Characteristic
+from capa.features.common import ARCH_X32, ARCH_X64, OS_WINDOWS, FORMAT_PE, String, Characteristic
 
 
 def test_rule_ctor():
@@ -961,3 +961,20 @@ def test_os_features():
     r = capa.rules.Rule.from_yaml(rule)
     children = list(r.statement.get_children())
     assert (Characteristic(OS_WINDOWS) in children) == True
+
+
+def test_format_features():
+    rule = textwrap.dedent(
+        """
+        rule:
+            meta:
+                name: test rule
+                scope: file
+            features:
+                - and:
+                    - characteristic: format/pe
+        """
+    )
+    r = capa.rules.Rule.from_yaml(rule)
+    children = list(r.statement.get_children())
+    assert (Characteristic(FORMAT_PE) in children) == True
\ No newline at end of file