Use cryptography for DH parameters

2 files changed, 40 insertions, 13 deletions

diff --git a/openid/dh.py b/openid/dh.py
index aeb00b5..37bb677 100644
--- a/openid/dh.py
+++ b/openid/dh.py
@@ -1,16 +1,12 @@
+""""Utilities for Diffie-Hellman key exchange."""
 from __future__ import unicode_literals
 
 import six
+from cryptography.hazmat.primitives.asymmetric.dh import DHParameterNumbers
 
 from openid import cryptutil
 from openid.constants import DEFAULT_DH_GENERATOR, DEFAULT_DH_MODULUS
 
-if six.PY2:
-    long_int = long
-else:
-    assert six.PY3
-    long_int = int
-
 
 def _xor(a_b):
     # Python 2 only
@@ -30,16 +26,37 @@ def strxor(x, y):
 
 
 class DiffieHellman(object):
+    """Utility for Diffie-Hellman key exchange."""
+
+    def __init__(self, modulus, generator):
+        """Create a new instance.
+
+        @type modulus: Union[six.integer_types]
+        @type generator: Union[six.integer_types]
+        """
+        self.parameter_numbers = DHParameterNumbers(modulus, generator)
+        self._setPrivate(cryptutil.randrange(1, modulus - 1))
 
     @classmethod
     def fromDefaults(cls):
+        """Create Diffie-Hellman with the default modulus and generator."""
         return cls(DEFAULT_DH_MODULUS, DEFAULT_DH_GENERATOR)
 
-    def __init__(self, modulus, generator):
-        self.modulus = long_int(modulus)
-        self.generator = long_int(generator)
+    @property
+    def modulus(self):
+        """Return the prime modulus value.
 
-        self._setPrivate(cryptutil.randrange(1, modulus - 1))
+        @rtype: Union[six.integer_types]
+        """
+        return self.parameter_numbers.p
+
+    @property
+    def generator(self):
+        """Return the generator value.
+
+        @rtype: Union[six.integer_types]
+        """
+        return self.parameter_numbers.g
 
     def _setPrivate(self, private):
         """This is here to make testing easier"""
diff --git a/openid/test/test_dh.py b/openid/test/test_dh.py
index ddd84f0..838e746 100644
--- a/openid/test/test_dh.py
+++ b/openid/test/test_dh.py
@@ -6,7 +6,8 @@ import unittest
 
 import six
 
-from openid.dh import DiffieHellman, long_int, strxor
+from openid.constants import DEFAULT_DH_GENERATOR, DEFAULT_DH_MODULUS
+from openid.dh import DiffieHellman, strxor
 
 
 class TestStrXor(unittest.TestCase):
@@ -52,6 +53,15 @@ class TestStrXor(unittest.TestCase):
 
 
 class TestDiffieHellman(unittest.TestCase):
+    """Test `DiffieHellman` class."""
+
+    def test_modulus(self):
+        dh = DiffieHellman.fromDefaults()
+        self.assertEqual(dh.modulus, DEFAULT_DH_MODULUS)
+
+    def test_generator(self):
+        dh = DiffieHellman.fromDefaults()
+        self.assertEqual(dh.generator, DEFAULT_DH_GENERATOR)
 
     def _test_dh(self):
         dh1 = DiffieHellman.fromDefaults()
@@ -72,8 +82,8 @@ class TestDiffieHellman(unittest.TestCase):
         try:
             for line in f:
                 parts = line.strip().split(' ')
-                dh._setPrivate(long_int(parts[0]))
+                dh._setPrivate(int(parts[0]))
 
-                assert dh.public == long_int(parts[1])
+                assert dh.public == int(parts[1])
         finally:
             f.close()