freeze: better type annotations for Address value

capa/features/freeze/__init__.py

@@ -12,7 +12,7 @@ See the License for the specific language governing permissions and limitations
 import zlib
 import logging
 from enum import Enum
-from typing import Any, List, Tuple
+from typing import Any, List, Tuple, Union
 
 from pydantic import Field, BaseModel
 
@@ -46,7 +46,7 @@ class AddressType(str, Enum):
 
 class Address(HashableModel):
     type: AddressType
-    value: Any
+    value: Union[int, Tuple[int, int], None]
 
     @classmethod
     def from_capa(cls, a: capa.features.address.Address) -> "Address":
@@ -79,19 +79,26 @@ class Address(HashableModel):
 
     def to_capa(self) -> capa.features.address.Address:
         if self.type is AddressType.ABSOLUTE:
+            assert isinstance(self.value, int)
             return capa.features.address.AbsoluteVirtualAddress(self.value)
 
         elif self.type is AddressType.RELATIVE:
+            assert isinstance(self.value, int)
             return capa.features.address.RelativeVirtualAddress(self.value)
 
         elif self.type is AddressType.FILE:
+            assert isinstance(self.value, int)
             return capa.features.address.FileOffsetAddress(self.value)
 
         elif self.type is AddressType.DN_TOKEN:
+            assert isinstance(self.value, int)
             return capa.features.address.DNTokenAddress(self.value)
 
         elif self.type is AddressType.DN_TOKEN_OFFSET:
+            assert isinstance(self.value, tuple)
             token, offset = self.value
+            assert isinstance(token, int)
+            assert isinstance(offset, int)
             return capa.features.address.DNTokenOffsetAddress(token, offset)
 
         elif self.type is AddressType.NO_ADDRESS:
@@ -108,7 +115,11 @@ class Address(HashableModel):
             return True
 
         else:
-            return self.value < other.value
+            assert self.type == other.type
+            # mypy doesn't realize we've proven that either
+            # both are ints, or both are tuples of ints.
+            # and both of these are comparable.
+            return self.value < other.value  # type: ignore
 
 
 class GlobalFeature(HashableModel):

capa/render/verbose.py

@@ -37,15 +37,22 @@ from capa.engine import MatchResults
 
 def format_address(address: frz.Address) -> str:
     if address.type == frz.AddressType.ABSOLUTE:
+        assert isinstance(address.value, int)
         return capa.helpers.hex(address.value)
     elif address.type == frz.AddressType.RELATIVE:
+        assert isinstance(address.value, int)
         return f"base address+{capa.helpers.hex(address.value)}"
     elif address.type == frz.AddressType.FILE:
+        assert isinstance(address.value, int)
         return f"file+{capa.helpers.hex(address.value)}"
     elif address.type == frz.AddressType.DN_TOKEN:
+        assert isinstance(address.value, int)
         return f"token({capa.helpers.hex(address.value)})"
     elif address.type == frz.AddressType.DN_TOKEN_OFFSET:
+        assert isinstance(address.value, tuple)
         token, offset = address.value
+        assert isinstance(token, int)
+        assert isinstance(offset, int)
         return f"token({capa.helpers.hex(token)})+{capa.helpers.hex(offset)}"
     elif address.type == frz.AddressType.NO_ADDRESS:
         return "global"