add support for sts regional interface endpoints

neptune-python-utils/neptune_python_utils/endpoints.py

@@ -60,8 +60,8 @@ def __init__(self, uri, querystring, headers):
         self.headers = headers
 
 class Endpoint:
-    
-    def __init__(self, protocol, neptune_endpoint, neptune_port, suffix, region, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False): 
+
+    def __init__(self, protocol, neptune_endpoint, neptune_port, suffix, region, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False, endpoint_url=None):
 
         self.protocol = protocol
         self.neptune_endpoint = neptune_endpoint
@@ -71,7 +71,8 @@ def __init__(self, protocol, neptune_endpoint, neptune_port, suffix, region, cre
         self.proxy_dns = proxy_dns
         self.proxy_port = proxy_port
         self.remove_host_header = remove_host_header
-
+        self.endpoint_url = endpoint_url
+
         if role_arn:
             self.role_arn = role_arn
             self.credentials = None
@@ -93,7 +94,10 @@ def _get_session_credentials(self):
     def _get_credentials(self):        
 
         if self.credentials is None:
-            sts = boto3.client('sts', region_name=self.region)
+            if self.endpoint_url:
+                sts = boto3.client('sts', region_name=self.region, endpoint_url=self.endpoint_url)
+            else:
+                sts = boto3.client('sts', region_name=self.region)
 
             role = sts.assume_role(
                 RoleArn=self.role_arn,
@@ -110,7 +114,10 @@ def _get_credentials(self):
         return self.credentials.get_frozen_credentials()
 
     def _new_credentials(self):
-        sts = boto3.client('sts', region_name=self.region)
+        if self.endpoint_url:
+            sts = boto3.client('sts', region_name=self.region, endpoint_url=self.endpoint_url)
+        else:
+            sts = boto3.client('sts', region_name=self.region)
 
         role = sts.assume_role(
             RoleArn=self.role_arn,
@@ -169,8 +176,8 @@ def get_headers():
 
 
 class Endpoints:
-    
-    def __init__(self, neptune_endpoint=None, neptune_port=None, region_name=None, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False):
+
+    def __init__(self, neptune_endpoint=None, neptune_port=None, region_name=None, credentials=None, role_arn=None, proxy_dns=None, proxy_port=8182, remove_host_header=False, endpoint_url=None):
 
         if neptune_endpoint is None:
             assert ('NEPTUNE_CLUSTER_ENDPOINT' in os.environ), 'neptune_endpoint is missing.'
@@ -188,12 +195,13 @@ def __init__(self, neptune_endpoint=None, neptune_port=None, region_name=None, c
         else:
             session = boto3.session.Session()
             self.region = session.region_name
-            
+
         self.credentials = credentials
         self.role_arn = role_arn
         self.proxy_dns = proxy_dns
         self.proxy_port = proxy_port
         self.remove_host_header = remove_host_header
+        self.endpoint_url = endpoint_url
 
 
     def gremlin_endpoint(self):
@@ -218,5 +226,4 @@ def sparql_stream_endpoint(self):
         return self.__endpoint('https', self.neptune_endpoint, self.neptune_port, 'sparql/stream')
 
     def __endpoint(self, protocol, neptune_endpoint, neptune_port, suffix):
-        return Endpoint(protocol, neptune_endpoint, neptune_port, suffix, self.region, self.credentials, self.role_arn, self.proxy_dns, self.proxy_port, self.remove_host_header)
-
+        return Endpoint(protocol, neptune_endpoint, neptune_port, suffix, self.region, self.credentials, self.role_arn, self.proxy_dns, self.proxy_port, self.remove_host_header, self.endpoint_url)

neptune-python-utils/readme.md

@@ -80,6 +80,20 @@ from neptune_python_utils.endpoints import Endpoints
 endpoints = Endpoints(role_arn='arn:aws:iam::...')
 ```
 
+To use it in a Secure Network Environment:
+
+Suppose your service handles critical PII data, so you have configured your VPC to not allow internet access by having no PublicSubnets and NAT Gateways in your VPC. Resources in such a VPC won't even have access to AWS services' endpoints. To access AWS services, you'd have to add InterfaceEndpoints or GatewayEndpoints of specific services in your VPC.   
+In such cases this client will render useless due to limitations of STS boto3 client unless you pass in the `endpoint_url` parameter.
+This parameter takes input the regional STS InterfaceEndpoint, which you can find here - [Using AWS STS interface VPC endpoints](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_sts_vpce.html)
+
+For more details, please visit this PR - [add support for sts regional interface endpoints](https://github.com/awslabs/amazon-neptune-tools/pull/227) 
+
+```
+from neptune_python_utils.endpoints import Endpoints
+
+endpoints = Endpoints(**other_kwargs, endpoint_url='https://sts.eu-west-1.amazonaws.com')
+```
+
 #### Proxies
 
 If you want to connect to Neptune via a proxy – a bastion host, [application load balancer or network load balancer](https://github.com/aws-samples/aws-dbs-refarch-graph/tree/master/src/connecting-using-a-load-balancer) – you must supply the proxy DNS and port to an `Endpoints` instance: