Remove duplicate code in equality check and include 0-padding

cpe_version.py

@@ -40,10 +40,11 @@ def __eq__(self, other):
         if (parts and not other_parts) or (not parts and other_parts):
             return False
 
-        # check for equality if one version has trailing ".0"
-        same_prefix = True
-        for part_idx in range(min(len(parts), len(other_parts))):
+        # iterate over pairs of parts for both versions
+        min_part_length = min(len(parts), len(other_parts))
+        for part_idx in range(min_part_length):
             part, other_part = parts[part_idx], other_parts[part_idx]
+
             # check if versions are in form 'update123' or 'u123' (e.g. cpe:2.3:a:trendmicro:deep_security_agent:20.0:update1559:)
             if part in ('u', 'update') and other_part in ('u', 'update'):
                 continue
@@ -52,27 +53,27 @@ def __eq__(self, other):
             if ONLY_ZEROES_STR_RE.match(part) and ONLY_ZEROES_STR_RE.match(other_part):
                 continue
 
-            if part.lower() != other_part.lower():
-                same_prefix = False
+            # right-pad with '0' to make both parts the same length
+            if len(part) < len(other_part) and part[0] in string.digits:
+                part = part.rjust(len(other_part), '0')
+            if len(other_part) < len(part) and other_part[0] in string.digits:
+                other_part = other_part.rjust(len(part), '0')
 
-        if same_prefix:
-            if len(parts) > len(other_parts) and all(not x or x == "0" for x in parts[part_idx+1:]):
-                return True
-            if len(other_parts) > len(parts) and all(not x or x == "0" for x in other_parts[part_idx+1:]):
-                return True
-            if len(parts) == len(other_parts):
-                for i in range(len(parts)):
-                    # check if versions are in form 'update123' or 'u123' (e.g. cpe:2.3:a:trendmicro:deep_security_agent:20.0:update1559:)
-                    if parts[i] in ('u', 'update') and other_parts[i] in ('u', 'update'):
-                        continue
-                    if parts[i] in ('p', 'patch') and other_parts[i] in ('p', 'patch'):
-                        continue
-                    if ONLY_ZEROES_STR_RE.match(parts[i]) and ONLY_ZEROES_STR_RE.match(other_parts[i]):
-                        continue
-                    if parts[i] != other_parts[i]:
-                        return False
-                return True
-        return False
+            # if any parts at any shared index are different --> versions not equal
+            if part != other_part:
+                return False
+
+        # check that either version is not a zero extended variant of the other, e.g. 1.5 and 1.5.0
+        if len(parts) > len(other_parts):
+            for part_idx in range(min_part_length, len(parts)):
+                if any(char not in ('.', '0') for char in parts[part_idx]):
+                    return False
+        if len(other_parts) > len(parts):
+            for part_idx in range(min_part_length, len(other_parts)):
+                if any(char not in ('.', '0') for char in other_parts[part_idx]):
+                    return False
+
+        return True
 
     def __gt__(self, other):
         return not (self <= other)

tests/test_version_comparison.py

@@ -30,6 +30,12 @@ def test_version_comparison(self):
         self.assertTrue(CPEVersion('1.3-a') < CPEVersion('1.3.0-rc1'))
         self.assertFalse(CPEVersion('1.3.1-4.5') < CPEVersion('1.3.0-4.5.6'))
         self.assertTrue(CPEVersion('20.0') < CPEVersion('20.0.0-3445'))
+        self.assertTrue(CPEVersion('1.2.3+01') == CPEVersion('1.2.3+1'))
+        self.assertTrue(CPEVersion('1.2.3+a.01') == CPEVersion('1.2.3+a.1'))
+        self.assertTrue(CPEVersion('1.2.3+a.000001') == CPEVersion('1.2.3+a.1'))
+        self.assertTrue(CPEVersion('3.1.4-l.o.n.g.e.r.rc.4') < CPEVersion('3.1.4-l.o.n.g.e.r.rc.5'))
+        self.assertTrue(CPEVersion('3.1.2147483647') < CPEVersion('3.1.2147483647-1'))
+        self.assertTrue(CPEVersion('3.1.2147483647-5') > CPEVersion('3.1.2147483647-1'))
 
 
 if __name__ == '__main__':