Create an HMAC key in AWS KMS

# Create an HMAC key in AWS KMS
# Begin by creating an HMAC key in AWS KMS. You can use the AWS KMS console or call the CreateKey API action. The following example shows creation of a 256-bit HMAC key:

import time
import json
import base64

import boto3

kms = boto3.client('kms')

# Use CreateKey API to create a 256-bit key for HMAC
key_id = kms.create_key(
    KeySpec='HMAC_256',
    KeyUsage='GENERATE_VERIFY_MAC'
)['KeyMetadata']['KeyId']

# ---------------------------------------------------------------------------------------

# How to protect HMACs inside AWS KMS > Create an HMAC key in AWS KMS
# https://aws.amazon.com/blogs/security/how-to-protect-hmacs-inside-aws-kms/


def base64_url_encode(data):
    return base64.b64encode(data, b'-_').rstrip(b'=')


# Payload contains simple claim and an issuance timestamp
payload = json.dumps({
    "does_kms_support_hmac": "yes",
    "iat": int(time.time())
}).encode("utf8")

# Header describes the algorithm and AWS KMS key ID to be used for signing
header = json.dumps({
    "typ": "JWT",
    "alg": "HS256",
    "kid": key_id  # This key_id is from the "Create an HMAC key in AWS KMS" #example. The "Verify the signed JWT" example will later #assert that the input header has the same value of the #key_id
}).encode("utf8")

# Message to sign is of form <header_b64>.<payload_b64>
message = base64_url_encode(header) + b'.' + base64_url_encode(payload)

# Generate MAC using GenerateMac API of AWS KMS
mac = kms.generate_mac(
    KeyId=key_id,  # This key_id is from the "Create an HMAC key in AWS KMS"
    # example
    MacAlgorithm='HMAC_SHA_256',
    Message=message
)['Mac']

# Form JWT token of form <header_b64>.<payload_b64>.<mac_b64>
jwt_token = message + b'.' + base64_url_encode(mac)


# ---------------------------------------------------------------------------------------

# Verify the signed JWT

def base64_url_decode(data):
    return base64.b64decode(data + b'=' * (4 - len(data) % 4), b'-_')


# Parse out encoded header, payload, and MAC from the token
message, mac_b64 = jwt_token.rsplit(b'.', 1)
header_b64, payload_b64 = message.rsplit(b'.', 1)

# Decode header and verify its contents match expectations
header_map = json.loads(base64_url_decode(header_b64).decode("utf8"))
assert header_map == {
    "typ": "JWT",
    "alg": "HS256",
    "kid": key_id  # This key_id is from the "Create an HMAC key in AWS KMS"
    # example
}

# Verify the MAC using VerifyMac API of AWS KMS. # If the verification fails, this will throw an error.
kms.verify_mac(
    KeyId=key_id,  # This key_id is from the "Create an HMAC key in AWS KMS"
    # example
    MacAlgorithm='HMAC_SHA_256',
    Message=message,
    Mac=base64_url_decode(mac_b64)
)

# Decode payload for use application-specific validation/processing
payload_map = json.loads(base64_url_decode(payload_b64).decode("utf8"))

# ---------------------------------------------------------------------------------------

# Create separate roles to control who has access to generate HMACs and who has access to validate HMACs

# {
#     "Id": "example-jwt-policy",
#     "Version": "2012-10-17",
#     "Statement": [
#         {
#             "Sid": "Allow use of the key for creating JWTs",
#             "Effect": "Allow",
#             "Principal": {
#                 "AWS": "arn:aws:iam::111122223333:role/JwtProducer"
#             },
#             "Action": [
#                 "kms:GenerateMac"
#             ],
#             "Resource": "*"
#         },
#         {
#             "Sid": "Allow use of the key for validating JWTs",
#             "Effect": "Allow",
#             "Principal": {
#                 "AWS": "arn:aws:iam::111122223333:role/JwtConsumer"
#             },
#             "Action": [
#                 "kms:VerifyMac"
#             ],
#             "Resource": "*"
#         }
#     ]
# }