handle_auth_code_callback2.py

def handle_auth_code_callback(body, event, auth_code, aad_object_id):
    # ...
    token_response = exchange_code_for_token(auth_code, TENANT_ID, CLIENT_ID, CLIENT_SECRET)
    
    # Extract the access token and expiration time
    access_token = token_response["access_token"]
    expires_in = token_response["expires_in"]
    
    # Calculate expiration time in seconds since epoch
    expires_at = int(time.time()) + expires_in
    
    ### Encrypt the access token using the CMK key
    
    # Initialize the KMS client
    kms = boto3.client('kms', region_name='us-east-1')  # Change region if needed

    # Encrypt the access token
    encrypted_token = kms.encrypt(
        KeyId=cmk_key_alias,
        Plaintext=access_token.encode("utf-8")
    )
    
    # Base64 encode the encrypted token
    encrypted_token_base64 = base64.b64encode(encrypted_token['CiphertextBlob']).decode('utf-8')