Better error messages when missing cryptography package

2 files changed, 48 insertions, 13 deletions

diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index 51e8f16..9aa50e7 100644
--- a/jwt/algorithms.py
+++ b/jwt/algorithms.py
@@ -25,6 +25,34 @@ except ImportError:
     has_crypto = False
 
 
+def _get_crypto_algorithms():
+    crypto_algorithms = {
+        'RS256': None,
+        'RS384': None,
+        'RS512': None,
+        'ES256': None,
+        'ES384': None,
+        'ES512': None,
+        'PS256': None,
+        'PS384': None,
+        'PS512': None
+    }
+
+    if has_crypto:
+        crypto_algorithms['RS256'] = RSAAlgorithm(RSAAlgorithm.SHA256)
+        crypto_algorithms['RS384'] = RSAAlgorithm(RSAAlgorithm.SHA384)
+        crypto_algorithms['RS512'] = RSAAlgorithm(RSAAlgorithm.SHA512)
+        crypto_algorithms['ES256'] = ECAlgorithm(ECAlgorithm.SHA256)
+        crypto_algorithms['ES384'] = ECAlgorithm(ECAlgorithm.SHA384)
+        crypto_algorithms['ES512'] = ECAlgorithm(ECAlgorithm.SHA512)
+        crypto_algorithms['PS256'] = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256)
+        crypto_algorithms['PS384'] = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA384)
+        crypto_algorithms['PS512'] = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA512)
+
+    return crypto_algorithms
+
+
+
 def get_default_algorithms():
     """
     Returns the algorithms that are implemented by the library.
@@ -37,21 +65,22 @@ def get_default_algorithms():
     }
 
     if has_crypto:
-        default_algorithms.update({
-            'RS256': RSAAlgorithm(RSAAlgorithm.SHA256),
-            'RS384': RSAAlgorithm(RSAAlgorithm.SHA384),
-            'RS512': RSAAlgorithm(RSAAlgorithm.SHA512),
-            'ES256': ECAlgorithm(ECAlgorithm.SHA256),
-            'ES384': ECAlgorithm(ECAlgorithm.SHA384),
-            'ES512': ECAlgorithm(ECAlgorithm.SHA512),
-            'PS256': RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256),
-            'PS384': RSAPSSAlgorithm(RSAPSSAlgorithm.SHA384),
-            'PS512': RSAPSSAlgorithm(RSAPSSAlgorithm.SHA512)
-        })
+        crypto_algorithms = _get_crypto_algorithms()
+        default_algorithms.update(crypto_algorithms)
 
     return default_algorithms
 
 
+def get_crypto_algorithms():
+    """
+    Returns a set of algorithm names that require the cryptography package to
+    be installed in order to use.
+    """
+    crypto_algorithms = _get_crypto_algorithms().keys()
+    return set(crypto_algorithms)
+
+
+
 class Algorithm(object):
     """
     The interface for an algorithm used to sign and verify tokens.
diff --git a/jwt/api_jws.py b/jwt/api_jws.py
index 177f5ff..b55734f 100644
--- a/jwt/api_jws.py
+++ b/jwt/api_jws.py
@@ -4,7 +4,9 @@ import warnings
 
 from collections import Mapping
 
-from .algorithms import Algorithm, get_default_algorithms  # NOQA
+from .algorithms import (
+    Algorithm, get_default_algorithms, has_crypto, get_crypto_algorithms  # NOQA
+)
 from .compat import binary_type, string_types, text_type
 from .exceptions import DecodeError, InvalidAlgorithmError, InvalidTokenError
 from .utils import base64url_decode, base64url_encode, merge_dict
@@ -99,7 +101,11 @@ class PyJWS(object):
             signature = alg_obj.sign(signing_input, key)
 
         except KeyError:
-            raise NotImplementedError('Algorithm not supported')
+            if not has_crypto and algorithm in get_crypto_algorithms():
+                raise NotImplementedError('"cryptography" package must be '
+                                          'installed to use this algorithm')
+            else:
+                raise NotImplementedError('Algorithm not supported')
 
         segments.append(base64url_encode(signature))