Verify public key and retrieve secret

verify_public_key_and_retrieve_secret.py

import json
import boto3
import base64
import os
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from botocore.exceptions import ClientError

def get_secret(secret_name):
    """Retrieve secret from AWS Secrets Manager."""
    client = boto3.client('secretsmanager')
    try:
        response = client.get_secret_value(SecretId=secret_name)
        return response['SecretString']
    except ClientError as e:
        print("Cannot retrieve the secret:", e)
        return None

def lambda_handler(event, context):
    # Get the hashed public key from the header
    client_hashed_key = event['headers'].get('X-Client-Public-Key-Hash')

    # Retrieve the private key from Secrets Manager
    private_key_data = get_secret(os.environ['PRIVATE_KEY_SECRET_ID'])
    private_key = serialization.load_pem_private_key(
        private_key_data.encode(),
        password=None,
    )

    # Generate the public key from the private key
    public_key = private_key.public_key()
    public_key_bytes = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )

    # Hash the public key
    digest = hashes.Hash(hashes.SHA256())
    digest.update(public_key_bytes)
    computed_hash = digest.finalize().hex()

    # Check if the provided hash matches the computed hash
    if client_hashed_key == computed_hash:
        # If matched, retrieve the secret string to return to the client
        secret_string = get_secret('SECRET_STRING_TO_SEND_BACK')
        return {
            'statusCode': 200,
            'body': json.dumps({'secret': secret_string})
        }
    else:
        # If not matched, return an error message
        return {
            'statusCode': 403,
            'body': json.dumps({'message': 'Unauthorized'})
        }