Merge pull request #7 from lyft/SEC-620-allow-list-of-validation-keys

Allow list of keys for user verification

Author: Stype

README.md

@@ -70,10 +70,10 @@ token = generator.get_token()
 ```python
 import kmsauth
 validator = kmsauth.KMSTokenValidator(
-    # KMS key to use for service authentication
-    'alias/authnz-production',
-    # KMS key to use for user authentication
-    'alias/authnz-users-production',
+    # KMS keys to use for service authentication
+    ['alias/authnz-production'],
+    # KMS keys to use for user authentication
+    ['alias/authnz-users-production', '6655d2a8-0606-4727-a1f6-f5b6a6754377'],
     # The context of this validation (the "to" context to validate against)
     'confidant-production',
     # Find the KMS keys in this region
@@ -88,10 +88,10 @@ context into the validator:
 ```python
 import kmsauth
 validator = kmsauth.KMSTokenValidator(
-    # KMS key to use for service authentication
-    'alias/authnz-production',
-    # KMS key to use for user authentication
-    'alias/authnz-users-production',
+    # KMS keys to use for service authentication
+    ['alias/authnz-production'],
+    # KMS keys to use for user authentication
+    ['alias/authnz-users-production', '6655d2a8-0606-4727-a1f6-f5b6a6754377'],
     # The context of this validation (the "to" context to validate against)
     'confidant-production',
     # Find the KMS keys in this region

kmsauth/__init__.py

@@ -7,6 +7,8 @@
 import sys
 import copy
 
+from botocore.vendored import six
+
 import kmsauth.services
 from kmsauth.utils import lru
 
@@ -35,9 +37,9 @@ def __init__(
         """Create a KMSTokenValidator object.
 
         Args:
-            auth_key: The KMS key ARN or alias to use for service
+            auth_key: A list of KMS key ARNs or aliases to use for service
                 authentication. Required.
-            user_auth_key: The KMS key ARN or alias to use for user
+            user_auth_key: A list of KMS key ARNs or aliases to use for user
                 authentication. Required.
             to_auth_context: The KMS encryption context to use for the to
                 context for authentication. Required.
@@ -103,6 +105,21 @@ def _validate(self):
                 'minimum_token_version can not be greater than'
                 ' self.minimum_token_version'
             )
+        self.auth_key = self._format_auth_key(self.auth_key)
+        self.user_auth_key = self._format_auth_key(self.user_auth_key)
+
+    def _format_auth_key(self, keys):
+        if isinstance(keys, six.string_types):
+            logging.debug(
+                'Passing auth key as string is deprecated, and will be removed'
+                ' in 1.0.0'
+            )
+            return [keys]
+        elif (keys is None or isinstance(keys, list)):
+            return keys
+        raise ConfigurationError(
+            'auth_key and user_auth_key must be a string, list, or None'
+        )
 
     def _get_key_arn(self, key):
         if key not in self.KEY_METADATA:
@@ -126,8 +143,9 @@ def _get_key_alias_from_cache(self, key_arn):
     def _valid_service_auth_key(self, key_arn):
         if self.auth_key is None:
             return False
-        if key_arn == self._get_key_arn(self.auth_key):
-            return True
+        for key in self.auth_key:
+            if key_arn == self._get_key_arn(key):
+                return True
         for key in self.scoped_auth_keys:
             if key_arn == self._get_key_arn(key):
                 return True
@@ -136,8 +154,9 @@ def _valid_service_auth_key(self, key_arn):
     def _valid_user_auth_key(self, key_arn):
         if self.user_auth_key is None:
             return False
-        if key_arn == self._get_key_arn(self.user_auth_key):
-            return True
+        for key in self.user_auth_key:
+            if key_arn == self._get_key_arn(key):
+                return True
         return False
 
     def _parse_username(self, username):

tests/unit/kmsauth/kmsauth_test.py

@@ -18,7 +18,7 @@ class KMSTokenValidatorTest(unittest.TestCase):
     def test_validate_config(self):
         with self.assertRaises(kmsauth.ConfigurationError):
             kmsauth.KMSTokenValidator(
-                'alias/authnz-unittest',
+                ['alias/authnz-unittest'],
                 None,
                 'kmsauth-unittest',
                 'us-east-1',
@@ -27,7 +27,7 @@ def test_validate_config(self):
             )
         with self.assertRaises(kmsauth.ConfigurationError):
             kmsauth.KMSTokenValidator(
-                'alias/authnz-unittest',
+                ['alias/authnz-unittest'],
                 None,
                 'kmsauth-unittest',
                 'us-east-1',
@@ -36,7 +36,7 @@ def test_validate_config(self):
             )
         with self.assertRaises(kmsauth.ConfigurationError):
             kmsauth.KMSTokenValidator(
-                'alias/authnz-unittest',
+                ['alias/authnz-unittest'],
                 None,
                 'kmsauth-unittest',
                 'us-east-1',
@@ -45,7 +45,7 @@ def test_validate_config(self):
             )
         with self.assertRaises(kmsauth.ConfigurationError):
             kmsauth.KMSTokenValidator(
-                'alias/authnz-unittest',
+                ['alias/authnz-unittest'],
                 None,
                 'kmsauth-unittest',
                 'us-east-1',
@@ -54,20 +54,34 @@ def test_validate_config(self):
             )
         with self.assertRaises(kmsauth.ConfigurationError):
             kmsauth.KMSTokenValidator(
-                'alias/authnz-unittest',
+                ['alias/authnz-unittest'],
                 None,
                 'kmsauth-unittest',
                 'us-east-1',
                 # minimum can't be greater than maximum
                 minimum_token_version=2,
                 maximum_token_version=1
             )
+        with self.assertRaises(kmsauth.ConfigurationError):
+            kmsauth.KMSTokenValidator(
+                # kms key must be string, list, or None
+                1234,
+                None,
+                'kmsauth-unittest',
+                'us-east-1',
+            )
         assert(kmsauth.KMSTokenValidator(
             'alias/authnz-unittest',
             None,
             'kmsauth-unittest',
             'us-east-1'
         ))
+        assert(kmsauth.KMSTokenValidator(
+            ['alias/authnz-unittest'],
+            None,
+            'kmsauth-unittest',
+            'us-east-1'
+        ))
 
     def test__get_key_arn(self):
         validator = kmsauth.KMSTokenValidator(