allow overriding endpoint_url (#11)

* allow overriding endpoint_url

this allows support for using a kms reverse proxy

* fix caching when specifying endpoint_url

Author: doy-stripe

kmsauth/__init__.py

@@ -34,7 +34,8 @@ def __init__(
             maximum_token_version=2,
             auth_token_max_lifetime=60,
             aws_creds=None,
-            extra_context=None
+            extra_context=None,
+            endpoint_url=None
             ):
         """Create a KMSTokenValidator object.
 
@@ -54,6 +55,8 @@ def __init__(
             aws_creds: A dict of AccessKeyId, SecretAccessKey, SessionToken.
                 Useful if you wish to pass in assumed role credentials or MFA
                 credentials. Default: None
+            endpoint_url: A URL to override the default endpoint used to access
+                the KMS service. Default: None
         """
         self.auth_key = auth_key
         self.user_auth_key = user_auth_key
@@ -73,12 +76,14 @@ def __init__(
                 region=self.region,
                 aws_access_key_id=self.aws_creds['AccessKeyId'],
                 aws_secret_access_key=self.aws_creds['SecretAccessKey'],
-                aws_session_token=self.aws_creds['SessionToken']
+                aws_session_token=self.aws_creds['SessionToken'],
+                endpoint_url=endpoint_url
             )
         else:
             self.kms_client = kmsauth.services.get_boto_client(
                 'kms',
-                region=self.region
+                region=self.region,
+                endpoint_url=endpoint_url
             )
         if extra_context is None:
             self.extra_context = {}
@@ -308,7 +313,8 @@ def __init__(
             token_version=2,
             token_cache_file=None,
             token_lifetime=10,
-            aws_creds=None
+            aws_creds=None,
+            endpoint_url=None
             ):
         """Create a KMSTokenGenerator object.
 
@@ -326,6 +332,8 @@ def __init__(
             aws_creds: A dict of AccessKeyId, SecretAccessKey, SessionToken.
                 Useful if you wish to pass in assumed role credentials or MFA
                 credentials. Default: None
+            endpoint_url: A URL to override the default endpoint used to access
+                the KMS service. Default: None
         """
         self.auth_key = auth_key
         if auth_context is None:
@@ -343,12 +351,14 @@ def __init__(
                 region=self.region,
                 aws_access_key_id=self.aws_creds['AccessKeyId'],
                 aws_secret_access_key=self.aws_creds['SecretAccessKey'],
-                aws_session_token=self.aws_creds['SessionToken']
+                aws_session_token=self.aws_creds['SessionToken'],
+                endpoint_url=endpoint_url
             )
         else:
             self.kms_client = kmsauth.services.get_boto_client(
                 'kms',
-                region=self.region
+                region=self.region,
+                endpoint_url=endpoint_url
             )
         self._validate()
 

kmsauth/services.py

@@ -12,10 +12,16 @@ def get_boto_client(
         region=None,
         aws_access_key_id=None,
         aws_secret_access_key=None,
-        aws_session_token=None
+        aws_session_token=None,
+        endpoint_url=None
         ):
     """Get a boto3 client connection."""
-    cache_key = '{0}:{1}:{2}'.format(client, region, aws_access_key_id)
+    cache_key = '{0}:{1}:{2}:{3}'.format(
+        client,
+        region,
+        aws_access_key_id,
+        endpoint_url or ''
+    )
     if not aws_session_token:
         if cache_key in CLIENT_CACHE:
             return CLIENT_CACHE[cache_key]
@@ -29,7 +35,10 @@ def get_boto_client(
         logging.error("Failed to get {0} client.".format(client))
         return None
 
-    CLIENT_CACHE[cache_key] = session.client(client)
+    CLIENT_CACHE[cache_key] = session.client(
+        client,
+        endpoint_url=endpoint_url
+    )
     return CLIENT_CACHE[cache_key]
 
 
@@ -38,10 +47,16 @@ def get_boto_resource(
         region=None,
         aws_access_key_id=None,
         aws_secret_access_key=None,
-        aws_session_token=None
+        aws_session_token=None,
+        endpoint_url=None
         ):
     """Get a boto resource connection."""
-    cache_key = '{0}:{1}:{2}'.format(resource, region, aws_access_key_id)
+    cache_key = '{0}:{1}:{2}:{3}'.format(
+        resource,
+        region,
+        aws_access_key_id,
+        endpoint_url or ''
+    )
     if not aws_session_token:
         if cache_key in RESOURCE_CACHE:
             return RESOURCE_CACHE[cache_key]
@@ -55,7 +70,10 @@ def get_boto_resource(
         logging.error("Failed to get {0} resource.".format(resource))
         return None
 
-    RESOURCE_CACHE[cache_key] = session.resource(resource)
+    RESOURCE_CACHE[cache_key] = session.resource(
+        resource,
+        endpoint_url=endpoint_url
+    )
     return RESOURCE_CACHE[cache_key]