Leveraging AWS signed requests for Identity Proof

import os
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import get_credentials
from botocore.endpoint import URLLib3Session
import boto3
import requests

session = boto3.Session()
credentials = session.get_credentials()
credentials = credentials.get_frozen_credentials()
signer = SigV4Auth(credentials, 'sts', 'us-east-1')

endpoint = 'api.example.com'
request = AWSRequest(
    method='POST', 
    url="https://sts.amazonaws.com/", 
    data="Action=GetCallerIdentity&Version=2011-06-15", 
    headers={
        'Host': 'sts.amazonaws.com',
        'Content-Type': 'application/x-www-form-urlencoded; charset=utf-8',
        'X-Audience': endpoint,
    }
)
signer.add_auth(request)

if 'dev' in os.environ.get('PY_ENV', ''):
    print(request.headers)
    print(request.body)

# To send to STS directly ....
# sender = URLLib3Session()
# response = sender.send(request.prepare())

# To send to our API, we use the request boto created with the 
# signature components
headers = dict(request.headers)
headers.pop('Host')
response = requests.post(f'https://{endpoint}/key', headers=headers, data=request.body)
print(response.status_code)
print(response.headers)
print(response.text)

import os
import json
import boto3
import http.client
import urllib.parse
import socket
import xml.etree.ElementTree as ET

def deny_response(message):
    return {
    	"policyDocument": {
    		"Version": "2012-10-17",
    		"Statement": [
    			{
    				"Action": "execute-api:Invoke",
    				"Effect": "Deny",
    				"Resource": "*",
    			}
    		]
    	},
    	"context": {
    	    'message': f'{message}'
    	}
    }

def handler(event, context):
    authorized_accounts = [a for a in os.environ.get('AUTHD_ACCTS', '').split(",") if a]
    authorized_roles = [r for r in os.environ.get('AUTHD_ROLES', '').split(",") if r]

    headers = event.get('headers', {})

    # Since API Gateway routes based on hostnames we know the host can be 
    # trusted. To be extra sure you could instead compare against
    # an environment variable of X_AUDIENCE 
    host = headers['host'] if 'host' in headers else headers.get('Host')
    
    # For the client's safety, expect this header
    if 'x-audience' not in headers:
        print("Did not supply x-audience header")
        return deny_response("A x-audience header is required")
    
    # ... and then assert the audience and the host are the same
    # This protects the client from host replay attacks
    if headers.get('x-audience') != host:
        print(f"Mismatch with host and x-audience header - audience={headers.get('x-audience')} host={host}" )
        return deny_response("Mismatch with host and x-audience header")

    keep = ['authorization', "x-amz-date", "x-amz-security-token", 'x-audience']
    keep = [k for k in headers if k.lower() in keep]
    
    authd = {k.lower(): headers[k] for k in keep if k in headers}
    
    if 'x-authorization' in headers:
        authd['authorization'] = headers.get('x-authorization')
    
    if 'authorization' not in authd:
        print("No authorization provided")
        return deny_response("No authorization header provided")

    if 'x-audience' not in authd['authorization']:
        print("The header x-audience was not included in the signature")
        return deny_response("The header x-audience was not included in the signature")
    
    params = urllib.parse.urlencode({ 
        'Action': "GetCallerIdentity", 
        'Version': "2011-06-15" 
    })

    authd['User-Agent'] = 'Boto3/1.10.28 Python/3.8.5 Darwin/18.7.0 Botocore/1.13.50'
    authd['Accept-Encoding'] = 'identity'
    authd["Content-Type"] = "application/x-www-form-urlencoded; charset=utf-8"
    authd['Content-Length'] = len(params)

    conn = http.client.HTTPSConnection("sts.amazonaws.com") 
    conn.request("POST", "/", params, authd)
    response = conn.getresponse()
    data = response.read()

    print(response.status, response.reason)
    print(response.getheaders())

    if "ExpiredToken" in str(data):
        print("Using expired tokens ....")
        return deny_response("The token provided is expired")

    if response.status != 200:
        print("Request was not successful with sts")
        print(data)
        return deny_response("Request STS failed")
    
    
    tree = ET.ElementTree(ET.fromstring(data))
    root = tree.getroot()
    nodes = [n.getchildren() for n in root if n.tag.split("}")[-1] == "GetCallerIdentityResult"].pop()
    
    if not nodes:
        print("No nodes?")
        return deny_response("STS response not validated")
    
    data = {n.tag.split("}")[-1]: n.text for n in nodes}
    
    if 'Arn' not in data or 'Account' not in data:
        print("The sts data is missing bits")
        return deny_response("STS response not validated")
    
    print("Authorized "  + json.dumps(data))
    parts = data['Arn'].split("/")
    data['Username'] = parts.pop()
    data['Role'] = parts[1] if len(parts) > 1 else None

    if not data.get('Account') or data.get('Account') not in authorized_accounts:
        return deny_response("That account is not authorized")

    if not data.get('Role') or data.get('Role') not in authorized_roles:
        return deny_response("That role is not authorized")

    resource = event.get('methodArn') if 'methodArn' in event else event.get('routeArn')

    return {
    	"principalId": data['Username'],
    	"policyDocument": {
    		"Version": "2012-10-17",
    		"Statement": [
    			{
    				"Action": "execute-api:Invoke",
    				"Effect": "Allow",
    				"Resource": resource,
    			}
    		]
    	},
    	"context": data,
    }

import re
import os
import json
import boto3
import string 
import random 
import base64

# You should use secrets manager to generate randomness but alas this serves as an example
# aws secretsmanager get-random-password --password-length 32 --require-each-included-type --exclude-characters '"'@/\"'"' | jq -r ".RandomPassword"'
def generate_random(size=48):
    return ''.join(random.choices(string.ascii_letters +
                             string.digits, k = size)) 


def unauthorized(actor=None):
    return {
        'statusCode': 400,
        'body': json.dumps({
            'message': 'You are not authorized for this',
            'actor': actor,
        }),
    }


def handler(event, context):
    # Verify this is called with an authorizer in front
    actor = event.get('requestContext', {}).get('authorizer', None)
    if actor is None:
        return unauthorized('No authorizer')
    
    # ... which provides us a username
    username = actor.get('Username')
    if not username:
        return unauthorized(actor)
    
    plan_id = os.environ['USAGE_PLAN']
    api_id = os.environ['API_ID']
    stages = [s.strip() for s in os.environ['API_STAGES'].split(',') if s.strip()]
    client = boto3.client('apigateway')
    already_has_key = False
    key_id = None
    try:
        response = client.get_api_keys(
            nameQuery=username,
            includeValues=False
        )
        items = response.get('items', [])
        already_has_key = len(items) > 0
        key_id = items[0]['id'] if already_has_key else None

    except Exception as err:
        print(err)
        err_string = re.sub(r"\d", "x", f'{err}')
        return {
            'statusCode': 500,
            'body': json.dumps({"error": err_string})
        }

    if already_has_key:
        return {
            'statusCode': 200,
            'body': json.dumps({
                "message": "You already generated an api key",
                'key_id': key_id,
                "username": username,
            })
        }

    key_value = username + ":" + generate_random()
    key_value = base64.b64encode(key_value.encode('ascii')).decode('ascii')

    try:    
        response = client.create_api_key(
            name=username, 
            description='Self service created key', 
            enabled=True, 
            value=key_value,
            stageKeys=[{'restApiId': api_id, 'stageName': s} for s in stages],
            tags={
                'SrcAccount': actor.get('Account'),
                'Role': actor.get('Role'),
            }
        )
    
        client.create_usage_plan_key(
            usagePlanId=plan_id,
            keyId=response['id'],
            keyType='API_KEY')    
    except Exception as err:
        print(err)
        err_string = re.sub(r"\d", "x", f'{err}')
        return {
            'statusCode': 500,
            'body': json.dumps({"error": err_string})
        }

    return {
        'statusCode': 200,
        'body': json.dumps({
            'username': username,
            'key_id': response['id'],
            'api_key': response['value']
        }),
    }