Follow the review advice: turn URI generator into methods.

5 files changed, 19 insertions, 36 deletions

diff --git a/src/cryptography/hazmat/primitives/twofactor/hotp.py b/src/cryptography/hazmat/primitives/twofactor/hotp.py
index ba228b40..f59f551c 100644
--- a/src/cryptography/hazmat/primitives/twofactor/hotp.py
+++ b/src/cryptography/hazmat/primitives/twofactor/hotp.py
@@ -15,6 +15,7 @@ from cryptography.hazmat.backends.interfaces import HMACBackend
 from cryptography.hazmat.primitives import constant_time, hmac
 from cryptography.hazmat.primitives.hashes import SHA1, SHA256, SHA512
 from cryptography.hazmat.primitives.twofactor import InvalidToken
+from cryptography.hazmat.primitives.twofactor.utils import generate_uri
 
 
 class HOTP(object):
@@ -59,3 +60,8 @@ class HOTP(object):
         offset = six.indexbytes(hmac_value, len(hmac_value) - 1) & 0b1111
         p = hmac_value[offset:offset + 4]
         return struct.unpack(">I", p)[0] & 0x7fffffff
+
+    def get_provisioning_uri(self, account_name, counter, issuer=None):
+        return generate_uri(self, 'hotp', account_name, issuer, [
+            ('counter', int(counter)),
+        ])
diff --git a/src/cryptography/hazmat/primitives/twofactor/totp.py b/src/cryptography/hazmat/primitives/twofactor/totp.py
index 03df9292..9c8eddad 100644
--- a/src/cryptography/hazmat/primitives/twofactor/totp.py
+++ b/src/cryptography/hazmat/primitives/twofactor/totp.py
@@ -11,6 +11,7 @@ from cryptography.hazmat.backends.interfaces import HMACBackend
 from cryptography.hazmat.primitives import constant_time
 from cryptography.hazmat.primitives.twofactor import InvalidToken
 from cryptography.hazmat.primitives.twofactor.hotp import HOTP
+from cryptography.hazmat.primitives.twofactor.utils import generate_uri
 
 
 class TOTP(object):
@@ -31,3 +32,8 @@ class TOTP(object):
     def verify(self, totp, time):
         if not constant_time.bytes_eq(self.generate(time), totp):
             raise InvalidToken("Supplied TOTP value does not match.")
+
+    def get_provisioning_uri(self, account_name, issuer=None):
+        return generate_uri(self._hotp, 'totp', account_name, issuer, [
+            ('period', int(self._time_step)),
+        ])
diff --git a/src/cryptography/hazmat/primitives/twofactor/utils.py b/src/cryptography/hazmat/primitives/twofactor/utils.py
index 43f50b30..89d38ff2 100644
--- a/src/cryptography/hazmat/primitives/twofactor/utils.py
+++ b/src/cryptography/hazmat/primitives/twofactor/utils.py
@@ -5,25 +5,7 @@ import base64
 from six.moves.urllib.parse import quote, urlencode
 
 
-__all__ = ['get_provisioning_uri']
-
-
-def get_provisioning_uri(otp, account_name, issuer=None, counter=None):
-    """Generates a provisioning URI which can be recognized by Two-Factor
-    Authentication Apps. See also: http://git.io/vkvvY
-
-    :param otp: An instance of
-                :class:`cryptography.hazmat.primitives.twofactor.hotp.HOTP` or
-                :class:`cryptography.hazmat.primitives.twofactor.totp.TOTP`.
-    :param account_name: The display name of account, such as
-                         ``'Alice Smith'`` or ``'alice@example.com'``.
-    :param issuer: The display name of issuer.
-    :param counter: The current value of counter. It is required for HOTP.
-    :return: The URI string.
-    :raises RuntimeError: if counter is missing but otp type is HOTP
-    """
-    hotp = getattr(otp, '_hotp', otp)
-
+def generate_uri(hotp, type_name, account_name, issuer, extra_parameters):
     parameters = [
         ('digits', hotp._length),
         ('secret', base64.b32encode(hotp._key)),
@@ -33,16 +15,10 @@ def get_provisioning_uri(otp, account_name, issuer=None, counter=None):
     if issuer is not None:
         parameters.append(('issuer', issuer))
 
-    if hotp is otp:
-        if counter is None:
-            raise RuntimeError('"counter" is required for HOTP')
-        parameters.append(('counter', int(counter)))
-
-    if hasattr(otp, '_time_step'):
-        parameters.append(('period', int(otp._time_step)))
+    parameters.extend(extra_parameters)
 
     uriparts = {
-        'type': otp.__class__.__name__.lower(),
+        'type': type_name,
         'label': ('%s:%s' % (quote(issuer), quote(account_name)) if issuer
                   else quote(account_name)),
         'parameters': urlencode(parameters),
diff --git a/tests/hazmat/primitives/twofactor/test_hotp.py b/tests/hazmat/primitives/twofactor/test_hotp.py
index ba40488a..3359dac2 100644
--- a/tests/hazmat/primitives/twofactor/test_hotp.py
+++ b/tests/hazmat/primitives/twofactor/test_hotp.py
@@ -14,7 +14,6 @@ from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.hashes import MD5, SHA1
 from cryptography.hazmat.primitives.twofactor import InvalidToken
 from cryptography.hazmat.primitives.twofactor.hotp import HOTP
-from cryptography.hazmat.primitives.twofactor.utils import get_provisioning_uri
 
 from ....utils import (
     load_nist_vectors, load_vectors_from_file, raises_unsupported_algorithm
@@ -97,18 +96,15 @@ class TestHOTP(object):
         secret = b"12345678901234567890"
         hotp = HOTP(secret, 6, SHA1(), backend)
 
-        assert get_provisioning_uri(hotp, "Alice Smith", counter=1) == (
+        assert hotp.get_provisioning_uri("Alice Smith", 1) == (
             "otpauth://hotp/Alice%20Smith?digits=6&secret=GEZDGNBV"
             "GY3TQOJQGEZDGNBVGY3TQOJQ&algorithm=SHA1&counter=1")
 
-        assert get_provisioning_uri(hotp, "Alice Smith", 'Foo', counter=1) == (
+        assert hotp.get_provisioning_uri("Alice Smith", 1, issuer='Foo') == (
             "otpauth://hotp/Foo:Alice%20Smith?digits=6&secret=GEZD"
             "GNBVGY3TQOJQGEZDGNBVGY3TQOJQ&algorithm=SHA1&issuer=Foo"
             "&counter=1")
 
-        with pytest.raises(RuntimeError):
-            get_provisioning_uri(hotp, "Alice Smith", 'World')  # counter lost
-
 
 def test_invalid_backend():
     secret = b"12345678901234567890"
diff --git a/tests/hazmat/primitives/twofactor/test_totp.py b/tests/hazmat/primitives/twofactor/test_totp.py
index 94c696f9..cd841ba6 100644
--- a/tests/hazmat/primitives/twofactor/test_totp.py
+++ b/tests/hazmat/primitives/twofactor/test_totp.py
@@ -11,7 +11,6 @@ from cryptography.hazmat.backends.interfaces import HMACBackend
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.twofactor import InvalidToken
 from cryptography.hazmat.primitives.twofactor.totp import TOTP
-from cryptography.hazmat.primitives.twofactor.utils import get_provisioning_uri
 
 from ....utils import (
     load_nist_vectors, load_vectors_from_file, raises_unsupported_algorithm
@@ -131,11 +130,11 @@ class TestTOTP(object):
         secret = b"12345678901234567890"
         totp = TOTP(secret, 6, hashes.SHA1(), 30, backend=backend)
 
-        assert get_provisioning_uri(totp, "Alice Smith") == (
+        assert totp.get_provisioning_uri("Alice Smith") == (
             "otpauth://totp/Alice%20Smith?digits=6&secret=GEZDGNBVG"
             "Y3TQOJQGEZDGNBVGY3TQOJQ&algorithm=SHA1&period=30")
 
-        assert get_provisioning_uri(totp, "Alice Smith", 'World') == (
+        assert totp.get_provisioning_uri("Alice Smith", 'World') == (
             "otpauth://totp/World:Alice%20Smith?digits=6&secret=GEZ"
             "DGNBVGY3TQOJQGEZDGNBVGY3TQOJQ&algorithm=SHA1&issuer=World"
             "&period=30")