vault-credential-provider.py

#!/usr/bin/env python

import os
from botocore.credentials import CredentialProvider, RefreshableCredentials
import requests
from datetime import datetime, timedelta

class VaultCredentialProvider(CredentialProvider):
    METHOD = 'vault'
    CANONICAL_NAME = 'VaultRole'

    def __init__(self, vault_url, vault_token):
        self._vault_url = vault_url
        self._vault_token = vault_token

    def load(self):
        metadata = self._load_from_vault()
        if not metadata:
            return None
        creds = RefreshableCredentials.create_from_metadata(
                metadata,
                method=self.METHOD,
                refresh_using=self._load_from_vault
                )
        return creds

    def _load_from_vault(self):
        try:
            now = datetime.utcnow()
            resp = requests.get(self._vault_url, headers={'X-Vault-Token':
                self._vault_token})
            resp.raise_for_status()
            vault_data = resp.json()
            expiry = now + timedelta(seconds=vault_data['lease_duration'])
            return {
                    'access_key': vault_data['data']['access_key'],
                    'secret_key': vault_data['data']['secret_key'],
                    'token': vault_data['data']['security_token'],
                    'expiry_time': expiry.isoformat() + 'Z'
                    }
        except:
            return None

def main():
    # see https://github.com/boto/boto3/issues/619#issuecomment-216980368
    from botocore.session import get_session
    from boto3.session import Session
    vault_addr = os.getenv('VAULT_ADDR')
    sts_path = "aws/sts/role"
    vault_cred_provider = VaultCredentialProvider("{0}/v1/{1}".format(vault_addr, sts_path),
            "c5b9220a-1804-9a66-acb2-0f7ffc73faca")
    bc_session = get_session()
    boto3_session = Session(botocore_session=bc_session)
    cred_provider = bc_session.get_component('credential_provider')
    cred_provider.insert_before('env', vault_cred_provider)
    sts_client = boto3_session.client('sts')
    print(sts_client.get_caller_identity())

if __name__ == "__main__":
    main()