From bfda997fdfe057646e176a95ec7d7d8a6ea1706c Mon Sep 17 00:00:00 2001
From: Willi Ballenthin <willi.ballenthin@gmail.com>
Date: Mon, 6 Jun 2022 15:55:12 -0600
Subject: [PATCH] freeze: support Class and Namespace features, too

---
 capa/features/freeze/features.py | 26 ++++++++++++++++++++++++++
 capa/render/vverbose.py          |  2 ++
 2 files changed, 28 insertions(+)

diff --git a/capa/features/freeze/features.py b/capa/features/freeze/features.py
index 3ddec4f3..8f8665ca 100644
--- a/capa/features/freeze/features.py
+++ b/capa/features/freeze/features.py
@@ -54,6 +54,12 @@ class FeatureModel(BaseModel):
         elif isinstance(self, StringFeature):
             return capa.features.common.String(self.string, description=self.description)
 
+        elif isinstance(self, ClassFeature):
+            return capa.features.common.Class(self.class_, description=self.description)
+
+        elif isinstance(self, NamespaceFeature):
+            return capa.features.common.Namespace(self.namespace, description=self.description)
+
         elif isinstance(self, BasicBlockFeature):
             return capa.features.basicblock.BasicBlock(description=self.description)
 
@@ -129,6 +135,12 @@ def feature_from_capa(f: capa.features.common.Feature) -> "Feature":
     elif isinstance(f, capa.features.common.String):
         return StringFeature(string=f.value, description=f.description)
 
+    elif isinstance(f, capa.features.common.Class):
+        return ClassFeature(class_=f.value, description=f.description)
+
+    elif isinstance(f, capa.features.common.Namespace):
+        return NamespaceFeature(namespace=f.value, description=f.description)
+
     elif isinstance(f, capa.features.basicblock.BasicBlock):
         return BasicBlockFeature(description=f.description)
 
@@ -231,6 +243,18 @@ class StringFeature(FeatureModel):
     description: Optional[str]
 
 
+class ClassFeature(FeatureModel):
+    type: str = "class"
+    class_: str = Field(alias="class")
+    description: Optional[str]
+
+
+class NamespaceFeature(FeatureModel):
+    type: str = "namespace"
+    namespace: str
+    description: Optional[str]
+
+
 class BasicBlockFeature(FeatureModel):
     type: str = "basic block"
     description: Optional[str]
@@ -293,6 +317,8 @@ Feature = Union[
     SubstringFeature,
     RegexFeature,
     StringFeature,
+    ClassFeature,
+    NamespaceFeature,
     APIFeature,
     NumberFeature,
     BytesFeature,
diff --git a/capa/render/vverbose.py b/capa/render/vverbose.py
index c4ab667f..b72504d9 100644
--- a/capa/render/vverbose.py
+++ b/capa/render/vverbose.py
@@ -131,6 +131,8 @@ def render_feature(ostream, match: rd.Match, feature: frzf.Feature, indent=0):
     if isinstance(feature, frzf.ImportFeature):
         # fixup access to Python reserved name
         value = feature.import_
+    if isinstance(feature, frzf.ClassFeature):
+        value = feature.class_
     else:
         value = getattr(feature, key)