Simple extending existing gremlin remote driver to work with AWS Neptune with IAM authentication

aws_gremlin_driver.py

import os, datetime, hashlib, hmac
from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from urllib.parse import urlparse


class AWSDriverRemoteConnection(DriverRemoteConnection):
    def __init__(self, url, traversal_source, protocol_factory=None,
                 transport_factory=None, pool_size=None, max_workers=None,
                 username="", password="", message_serializer=None,
                 graphson_reader=None, graphson_writer=None,
                 headers={}, aws_profile=None):

        DriverRemoteConnection.__init__(self, url, traversal_source, protocol_factory,
                                        transport_factory, pool_size, max_workers,
                                        username, password, message_serializer, graphson_reader,
                                        graphson_writer,
                                        do_headers(url, headers, aws_profile))



# Key derivation functions. See:
# http://docs.aws.amazon.com/general/latest/gr/signature-v4-examples.html#signature-v4-examples-python
def sign(key, msg):
    return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()


def getSignatureKey(key, dateStamp, regionName, serviceName):
    kDate = sign(('AWS4' + key).encode('utf-8'), dateStamp)
    kRegion = sign(kDate, regionName)
    kService = sign(kRegion, serviceName)
    kSigning = sign(kService, 'aws4_request')
    return kSigning


def get_profile_credentials(profile_name):
    print('Getting credentials for profile "{}"'.format(profile_name))
    from os import path
    from botocore.configloader import load_config
    config_file = path.join(path.expanduser("~"), '.aws/credentials')
    configs = load_config(config_file)

    if profile_name in configs.keys():
        config = configs[profile_name]
        os.environ['SERVICE_REGION'] = config['region']

        for k in ('aws_access_key_id', 'aws_secret_access_key', 'aws_session_token'):
            os.environ[k.upper()] = config[k]

        return config


def do_headers(url: str, headers: dict, aws_profile: str = None):
    if aws_profile is None and 'AWS_PROFILE' in os.environ.keys():
        aws_profile = os.environ['AWS_PROFILE']

    if aws_profile:
        get_profile_credentials(aws_profile)

    o = urlparse(url)
    # ************* REQUEST VALUES *************
    method = 'GET'
    service = 'neptune-db'
    host = o.netloc
    region = 'us-east-1'

    # Read AWS access key from env. variables or configuration file. Best practice is NOT
    # to embed credentials in code.
    session_token = os.getenv('AWS_SESSION_TOKEN', '')
    access_key = os.environ.get('AWS_ACCESS_KEY_ID')
    secret_key = os.environ.get('AWS_SECRET_ACCESS_KEY')
    if access_key is None or secret_key is None:
        print('No access key is available.')
        return headers
    # Create a date for headers and the credential string
    t = datetime.datetime.utcnow()
    amzdate = t.strftime('%Y%m%dT%H%M%SZ')
    datestamp = t.strftime('%Y%m%d')  # Date w/o time, used in credential scope

    # ************* TASK 1: CREATE A CANONICAL REQUEST *************
    # http://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html

    # Step 1 is to define the verb (GET, POST, etc.)--already done.

    # Step 2: Create canonical URI--the part of the URI from domain to query
    # string (use '/' if no path)
    canonical_uri = o.path

    # Step 3: Create the canonical query string. In this example (a GET request),
    # request parameters are in the query string. Query string values must
    # be URL-encoded (space=%20). The parameters must be sorted by name.
    # For this example, the query string is pre-formatted in the request_parameters variable.
    canonical_querystring = ''

    # Step 4: Create the canonical headers and signed headers. Header names
    # must be trimmed and lowercase, and sorted in code point order from
    # low to high. Note that there is a trailing \n.
    canonical_headers = 'host:' + host + '\n' + 'x-amz-date:' + amzdate + '\n'

    # Step 5: Create the list of signed headers. This lists the headers
    # in the canonical_headers list, delimited with ";" and in alpha order.
    # Note: The request can include any headers; canonical_headers and
    # signed_headers lists those that you want to be included in the
    # hash of the request. "Host" and "x-amz-date" are always required.
    signed_headers = 'host;x-amz-date'

    # Step 6: Create payload hash (hash of the request body content). For GET
    # requests, the payload is an empty string ("").
    payload_hash = hashlib.sha256(('').encode('utf-8')).hexdigest()

    # Step 7: Combine elements to create canonical request
    canonical_request = method + '\n' + canonical_uri + '\n' + canonical_querystring + '\n' + canonical_headers + '\n' + signed_headers + '\n' + payload_hash


    # ************* TASK 2: CREATE THE STRING TO SIGN*************
    # Match the algorithm to the hashing algorithm you use, either SHA-1 or
    # SHA-256 (recommended)
    algorithm = 'AWS4-HMAC-SHA256'
    credential_scope = datestamp + '/' + region + '/' + service + '/' + 'aws4_request'
    string_to_sign = algorithm + '\n' + amzdate + '\n' + credential_scope + '\n' + hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()

    # ************* TASK 3: CALCULATE THE SIGNATURE *************
    # Create the signing key using the function defined above.
    signing_key = getSignatureKey(secret_key, datestamp, region, service)

    # Sign the string_to_sign using the signing_key
    signature = hmac.new(signing_key, (string_to_sign).encode('utf-8'), hashlib.sha256).hexdigest()


    # ************* TASK 4: ADD SIGNING INFORMATION TO THE REQUEST *************
    # The signing information can be either in a query string value or in
    # a header named Authorization. This code shows how to use a header.
    # Create authorization header and add to request headers
    authorization_header = algorithm + ' ' + 'Credential=' + access_key + '/' + credential_scope + ', ' +  'SignedHeaders=' + signed_headers + ', ' + 'Signature=' + signature

    # The request can include any headers, but MUST include "host", "x-amz-date",
    # and (for this scenario) "Authorization". "host" and "x-amz-date" must
    # be included in the canonical_headers and signed_headers, as noted
    # earlier. Order here is not significant.
    # Python note: The 'host' header is added automatically by the Python 'requests' library.
    headers['x-amz-date'] = amzdate
    headers['Authorization'] = authorization_header

    if session_token:
        headers['x-amz-security-token'] = session_token

    return headers

Most of this is taken directly from this AWS resource.