fix: DNTokenOffsetAddress.__eq__ lacks type guard

Closes #3022

CHANGELOG.md

@@ -30,6 +30,8 @@
 
 - fix: address classes __eq__ and __lt__ assert on type @williballenthin #3021
 
+- fix: DNTokenOffsetAddress.__eq__ lacks type guard @williballenthin #3022
+
 ### capa Explorer Web
 
 ### capa Explorer IDA Pro plugin

capa/features/address.py

@@ -174,9 +174,13 @@ class DNTokenOffsetAddress(Address):
         self.offset = offset
 
     def __eq__(self, other):
+        if not isinstance(other, DNTokenOffsetAddress):
+            return NotImplemented
         return (self.token, self.offset) == (other.token, other.offset)
 
     def __lt__(self, other):
+        if not isinstance(other, DNTokenOffsetAddress):
+            return NotImplemented
         return (self.token, self.offset) < (other.token, other.offset)
 
     def __hash__(self):

tests/test_engine.py

@@ -17,7 +17,7 @@ import pytest
 import capa.features.address
 from capa.engine import Or, And, Not, Some, Range
 from capa.features.insn import Number
-from capa.features.address import ThreadAddress, ProcessAddress, DynamicCallAddress
+from capa.features.address import ThreadAddress, ProcessAddress, DynamicCallAddress, DNTokenOffsetAddress, AbsoluteVirtualAddress
 
 ADDR1 = capa.features.address.AbsoluteVirtualAddress(0x401001)
 ADDR2 = capa.features.address.AbsoluteVirtualAddress(0x401002)
@@ -49,6 +49,23 @@ def test_no_address_hash():
     assert d[addr_zero] == "zero"
 
 
+def test_dn_token_offset_address_cross_type_eq():
+    addr = DNTokenOffsetAddress(0x1000, 0x10)
+    assert (addr == AbsoluteVirtualAddress(0x1010)) is False
+    assert (addr == "not an address") is False
+    assert (addr == None) is False  # noqa: E711
+    assert (addr == DNTokenOffsetAddress(0x1000, 0x10)) is True
+    assert (addr == DNTokenOffsetAddress(0x1000, 0x11)) is False
+
+
+def test_dn_token_offset_address_cross_type_lt():
+    addr = DNTokenOffsetAddress(0x1000, 0x10)
+    assert addr.__lt__(AbsoluteVirtualAddress(0x1010)) is NotImplemented
+    assert addr.__lt__("not an address") is NotImplemented
+    assert (addr < DNTokenOffsetAddress(0x1000, 0x11)) is True
+    assert (addr < DNTokenOffsetAddress(0x1000, 0x10)) is False
+
+
 def test_number():
     assert bool(Number(1).evaluate({Number(0): {ADDR1}})) is False
     assert bool(Number(1).evaluate({Number(1): {ADDR1}})) is True