feat(common/parameter)!: Implement enum item getter by parameter numbers

c0r0n3r · Dec 28, 2023 · 72c3c08 · 72c3c08
1 parent 6f10aef
commit 72c3c08
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 2 deletions.
diff --git a/cryptodatahub/common/parameter.py b/cryptodatahub/common/parameter.py
@@ -77,4 +77,12 @@ def __str__(self):
         return '{}-bit {} {} DH parameter'.format(self.key_size, self.source.value.name, self.name)
 
 
-DHParamWellKnown = CryptoDataEnumBase('DHParamWellKnown', CryptoDataEnumBase.get_json_records(DHParamWellKnownParams))
+class DHParamWellKnownBase(CryptoDataEnumBase):
+    @classmethod
+    def from_parameter_numbers(cls, parameter_numbers):
+        return cls._from_attr('parameter_numbers', parameter_numbers)
+
+
+DHParamWellKnown = DHParamWellKnownBase(
+    'DHParamWellKnown', DHParamWellKnownBase.get_json_records(DHParamWellKnownParams)
+)
diff --git a/cryptodatahub/common/types.py b/cryptodatahub/common/types.py
@@ -502,7 +502,7 @@ def get_json(cls, param_class):
     @classmethod
     def _from_attr(cls, attr_name, value):
         for item in cls:
-            if getattr(item.value, attr_name) == value:
+            if value == getattr(item.value, attr_name):
                 return item
 
         raise InvalidValue(value, cls, attr_name)

diff --git a/test/common/test_parameter.py b/test/common/test_parameter.py
@@ -24,6 +24,13 @@ class TestDHParamWellKnown(TestClasses.TestJsonBase):
     def _get_class(cls):
         return DHParamWellKnown
 
+    def test_from_parameter_numbers(self):
+        dh_param = DHParamWellKnown.RFC2539_1024_BIT_MODP_GROUP
+        self.assertEqual(
+           DHParamWellKnown.from_parameter_numbers(dh_param.value.parameter_numbers),
+           dh_param
+        )
+
     def test_str(self):
         self.assertEqual(
             str(DHParamWellKnown.RFC2539_768_BIT_MODP_GROUP.value),