Add "generate_key_uri" utility for HOTP/TOTP.

3 files changed, 81 insertions, 0 deletions

diff --git a/src/cryptography/hazmat/primitives/twofactor/utils.py b/src/cryptography/hazmat/primitives/twofactor/utils.py
new file mode 100644
index 00000000..43f50b30
--- /dev/null
+++ b/src/cryptography/hazmat/primitives/twofactor/utils.py
@@ -0,0 +1,50 @@
+from __future__ import unicode_literals
+
+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)
+
+    parameters = [
+        ('digits', hotp._length),
+        ('secret', base64.b32encode(hotp._key)),
+        ('algorithm', hotp._algorithm.name.upper()),
+    ]
+
+    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)))
+
+    uriparts = {
+        'type': otp.__class__.__name__.lower(),
+        'label': ('%s:%s' % (quote(issuer), quote(account_name)) if issuer
+                  else quote(account_name)),
+        'parameters': urlencode(parameters),
+    }
+    return 'otpauth://{type}/{label}?{parameters}'.format(**uriparts)
diff --git a/tests/hazmat/primitives/twofactor/test_hotp.py b/tests/hazmat/primitives/twofactor/test_hotp.py
index a5d1c284..ba40488a 100644
--- a/tests/hazmat/primitives/twofactor/test_hotp.py
+++ b/tests/hazmat/primitives/twofactor/test_hotp.py
@@ -14,6 +14,7 @@ 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
@@ -92,6 +93,22 @@ class TestHOTP(object):
         with pytest.raises(TypeError):
             HOTP(secret, b"foo", SHA1(), backend)
 
+    def test_get_provisioning_uri(self, backend):
+        secret = b"12345678901234567890"
+        hotp = HOTP(secret, 6, SHA1(), backend)
+
+        assert get_provisioning_uri(hotp, "Alice Smith", counter=1) == (
+            "otpauth://hotp/Alice%20Smith?digits=6&secret=GEZDGNBV"
+            "GY3TQOJQGEZDGNBVGY3TQOJQ&algorithm=SHA1&counter=1")
+
+        assert get_provisioning_uri(hotp, "Alice Smith", 'Foo', counter=1) == (
+            "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 6039983e..94c696f9 100644
--- a/tests/hazmat/primitives/twofactor/test_totp.py
+++ b/tests/hazmat/primitives/twofactor/test_totp.py
@@ -11,6 +11,7 @@ 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
@@ -126,6 +127,19 @@ class TestTOTP(object):
 
         assert totp.generate(time) == b"94287082"
 
+    def test_get_provisioning_uri(self, backend):
+        secret = b"12345678901234567890"
+        totp = TOTP(secret, 6, hashes.SHA1(), 30, backend=backend)
+
+        assert get_provisioning_uri(totp, "Alice Smith") == (
+            "otpauth://totp/Alice%20Smith?digits=6&secret=GEZDGNBVG"
+            "Y3TQOJQGEZDGNBVGY3TQOJQ&algorithm=SHA1&period=30")
+
+        assert get_provisioning_uri(totp, "Alice Smith", 'World') == (
+            "otpauth://totp/World:Alice%20Smith?digits=6&secret=GEZ"
+            "DGNBVGY3TQOJQGEZDGNBVGY3TQOJQ&algorithm=SHA1&issuer=World"
+            "&period=30")
+
 
 def test_invalid_backend():
     secret = b"12345678901234567890"