@alliander-opensource/aws-jwt-sts 0.2.6

package/package.json

@@ -0,0 +1,56 @@
+{
+  "name": "@alliander-opensource/aws-jwt-sts",
+  "version": "0.2.6",
+  "author": {
+    "name": "Alliander NV"
+  },
+  "main": "dist/index.js",
+  "files": [
+    "dist/*",
+    "src/*"
+  ],
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/alliander-opensource/aws-jwt-sts.git"
+  },
+  "scripts": {
+    "build": "rm -rf dist && tsc",
+    "watch": "tsc -w",
+    "test": "jest --coverage",
+    "cdk": "cdk",
+    "lint": "eslint . --ext .ts"
+  },
+  "devDependencies": {
+    "@types/jest": "^29.5.0",
+    "@types/node": "18.15.11",
+    "@types/prettier": "2.7.2",
+    "@typescript-eslint/eslint-plugin": "^5.58.0",
+    "@typescript-eslint/parser": "^5.58.0",
+    "aws-cdk": "2.74.0",
+    "aws-sdk-client-mock": "^2.1.1",
+    "esbuild": "^0.17.17",
+    "eslint-config-standard": "^17.0.0",
+    "eslint-plugin-import": "^2.27.5",
+    "eslint-plugin-jest": "^27.2.1",
+    "eslint-plugin-n": "^15.7.0",
+    "eslint-plugin-promise": "^6.1.1",
+    "jest": "^29.5.0",
+    "jwt-decode": "^3.1.2",
+    "ts-jest": "^29.1.0",
+    "ts-node": "^10.9.1",
+    "typescript": "~5.0.4"
+  },
+  "dependencies": {
+    "@aws-lambda-powertools/logger": "^1.8.0",
+    "@aws-sdk/client-kms": "^3.312.0",
+    "@aws-sdk/client-s3": "^3.312.0",
+    "@types/aws-lambda": "^8.10.114",
+    "@types/jsrsasign": "^10.5.8",
+    "aws-cdk-lib": "2.74.0",
+    "base64url": "^3.0.1",
+    "constructs": "^10.1.313",
+    "jsrsasign": "^10.8.2",
+    "source-map-support": "^0.5.21"
+  }
+}

package/src/index.keyrotate.ts

@@ -0,0 +1,228 @@
+// SPDX-FileCopyrightText: 2023 Alliander NV
+//
+// SPDX-License-Identifier: Apache-2.0
+
+import {
+  KMSClient,
+  DescribeKeyCommand,
+  GetPublicKeyCommand,
+  CreateKeyCommand,
+  UpdateAliasCommand,
+  ScheduleKeyDeletionCommand,
+  TagResourceCommand,
+  Tag,
+  NotFoundException,
+  CreateAliasCommand
+} from '@aws-sdk/client-kms'
+import { S3Client, PutObjectCommand } from '@aws-sdk/client-s3'
+import { KEYUTIL, KJUR } from 'jsrsasign'
+
+const client = new KMSClient({})
+
+const ALIAS_PREVIOUS = 'alias/sts/PREVIOUS'
+const ALIAS_CURRENT = 'alias/sts/CURRENT'
+const ALIAS_PENDING = 'alias/sts/PENDING'
+
+const ALIASES: string[] = [
+  ALIAS_PREVIOUS,
+  ALIAS_CURRENT,
+  ALIAS_PENDING
+]
+
+export const handler = async (event: any): Promise<any> => {
+  // retrieve the step from the event
+  const step = (event.step)
+
+  // match the step with the corresponding function
+  switch (step) {
+    case 'deletePrevious':
+      await deletePrevious()
+      break
+    case 'movePrevious':
+      await movePrevious()
+      break
+    case 'moveCurrent':
+      await moveCurrent()
+      break
+    case 'createPending':
+      await createPending()
+      break
+    case 'generateArtifacts':
+      await generateJWKS()
+      await generateOpenIDConfiguration()
+      break
+
+    default:
+      console.log('invalid step')
+  }
+}
+
+async function deletePrevious () {
+  console.log('Deleting PREVIOUS aliased key')
+
+  const prevKeyId = await getKeyIdForAlias(ALIAS_PREVIOUS)
+  if (prevKeyId) {
+    const ScheduleDeleteResponse = await client.send(
+      new ScheduleKeyDeletionCommand({ KeyId: prevKeyId })
+    )
+    console.log(ScheduleDeleteResponse)
+  } else {
+    console.log('No PREVIOUS key at the moment, skip deletion')
+  }
+}
+
+async function movePrevious () {
+  console.log('moving PREVIOUS alias')
+  const currentKeyId = await getKeyIdForAlias(ALIAS_CURRENT)
+  if (currentKeyId) {
+    await updateOrCreateAlias(ALIAS_PREVIOUS, currentKeyId)
+  } else {
+    console.log('No CURRENT key at the moment, skip assigning the PREVIOUS alias to this key.')
+  }
+}
+
+async function moveCurrent () {
+  console.log('Moving CURRENT alias')
+
+  const pendingKeyId = await getKeyIdForAlias(ALIAS_PENDING)
+  if (pendingKeyId) {
+    await updateOrCreateAlias(ALIAS_CURRENT, pendingKeyId)
+  } else {
+    console.log('No PENDING key at the moment, skip assigning the CURRENT alias to this key.')
+  }
+}
+
+async function createPending () {
+  console.log('Creating new key for PENDING')
+
+  // Create new key
+  const createResponse = await client.send(new CreateKeyCommand({
+    KeySpec: 'RSA_2048',
+    KeyUsage: 'SIGN_VERIFY'
+  }))
+  console.log(createResponse)
+
+  // Update the new key with pending alias
+  await updateOrCreateAlias(ALIAS_PENDING, createResponse.KeyMetadata!.KeyId!)
+}
+
+async function updateOrCreateAlias (aliasName: string, keyId: string) {
+  try {
+    const updateResponse = await client.send(new UpdateAliasCommand({
+      AliasName: aliasName,
+      TargetKeyId: keyId
+    }))
+    console.log(updateResponse)
+  } catch (err) {
+    if (err instanceof NotFoundException) {
+      console.log('ALIAS not found, creating it.')
+      const createResponse = await client.send(new CreateAliasCommand({
+        AliasName: aliasName,
+        TargetKeyId: keyId
+      }))
+      console.log(createResponse)
+    } else {
+      throw (err)
+    }
+  }
+}
+
+async function getKeyIdForAlias (keyId: string) {
+  try {
+    const response = await client.send(new DescribeKeyCommand({ KeyId: keyId }))
+    console.log(response)
+    return response.KeyMetadata?.KeyId
+  } catch (err) {
+    if (err instanceof NotFoundException) {
+      return null
+    } else {
+      throw err
+    }
+  }
+}
+
+async function generateJWKS () {
+  const allKeys: object[] = []
+
+  for (const keyAlias of ALIASES) {
+    const keyId = await getKeyIdForAlias(keyAlias)
+    if (keyId) {
+      const jwkContents = await generateJWK(keyAlias)
+      await setKMSKeyTags(keyId, [{ TagKey: 'jwk_kid', TagValue: jwkContents.kid }])
+      allKeys.push(jwkContents)
+    }
+  }
+
+  const result = { keys: allKeys }
+
+  await uploadToS3('discovery/keys', result)
+}
+
+async function generateOpenIDConfiguration () {
+  const issuer = process.env.ISSUER
+
+  const openIdConfiguration = {
+    issuer,
+    jwks_uri: `${issuer}/discovery/keys`,
+    response_types_supported: [
+      'token'
+    ],
+    id_token_signing_alg_values_supported: [
+      'RS256'
+    ],
+    scopes_supported: [
+      'openid'
+    ],
+    token_endpoint_auth_methods_supported: [
+      'client_secret_basic'
+    ],
+    claims_supported: [
+      'aud',
+      'exp',
+      'iat',
+      'iss',
+      'sub'
+    ]
+  }
+
+  await uploadToS3('.well-known/openid-configuration', openIdConfiguration)
+}
+
+async function generateJWK (keyAlias: string): Promise<any> {
+  // Get the public key from kms
+  const getPubKeyResponse = await client.send(new GetPublicKeyCommand({ KeyId: keyAlias }))
+
+  // generate HEX format from the response (DER)
+  const pubKeyHex = Buffer.from(getPubKeyResponse.PublicKey as Uint8Array).toString('hex')
+
+  // Get the pub key in pem format
+  const pubKeyPem = KJUR.asn1.ASN1Util.getPEMStringFromHex(pubKeyHex, 'PUBLIC KEY')
+
+  // return the JWK format for the key
+  const jwk: any = KEYUTIL.getJWK(pubKeyPem)
+
+  jwk.use = 'sig'
+  jwk.alg = 'RS256'
+
+  return jwk
+}
+
+async function setKMSKeyTags (keyAlias: string, tags: Tag[]) {
+  return await client.send(new TagResourceCommand({ KeyId: keyAlias, Tags: tags }))
+}
+
+async function uploadToS3 (key: string, contents: object) {
+  // get S3 bucket from environment variables
+  const s3Bucket = process.env.S3_BUCKET
+
+  const s3client = new S3Client({})
+
+  // Write jwk to s3 bucket
+  await s3client.send(new PutObjectCommand({
+    Bucket: s3Bucket,
+    Key: key, // File name you want to save as in S3
+    Body: Buffer.from(JSON.stringify(contents)),
+    ContentType: 'application/json',
+    ContentEncoding: ''
+  }))
+}

package/src/index.sign.ts

@@ -0,0 +1,145 @@
+// SPDX-FileCopyrightText: 2023 Alliander NV
+//
+// SPDX-License-Identifier: Apache-2.0
+
+import { Context, APIGatewayProxyResult, APIGatewayEvent } from 'aws-lambda'
+import { KMSClient, SignCommand, DescribeKeyCommand, ListResourceTagsCommand, Tag } from '@aws-sdk/client-kms'
+import base64url from 'base64url'
+
+import { Logger } from '@aws-lambda-powertools/logger'
+
+const KEY_ALIAS_CURRENT = 'alias/sts/CURRENT'
+const logger = new Logger()
+
+export const handler = async (apiEvent: APIGatewayEvent, context: Context): Promise<APIGatewayProxyResult> => {
+  const identityArn = getARNFromIdentity(apiEvent.requestContext.identity?.userArn)
+  logger.debug(identityArn!)
+
+  if (identityArn === undefined || identityArn === null) {
+    logger.info(`Unable to resolve identityArn for userArn: ${apiEvent.requestContext.identity?.userArn}`)
+    return respond('Unable to resolve identity', 400)
+  }
+
+  let aud = process.env.DEFAULT_AUDIENCE
+
+  if (apiEvent.queryStringParameters && apiEvent.queryStringParameters.aud) {
+    aud = apiEvent.queryStringParameters.aud
+  }
+
+  const kms = new KMSClient({})
+
+  // Get KeyID which will be sent as kid in JWT token
+  const currentResponse = await kms.send(new DescribeKeyCommand({ KeyId: `${KEY_ALIAS_CURRENT}` }))
+  const currentKeyId = currentResponse.KeyMetadata?.KeyId
+
+  if (currentKeyId === undefined) {
+    return respond('KMS key could not be retrieved', 500)
+  }
+
+  // Retrieve Tags for KMS Key - the key is tagged with the `kid` from the JWK which is used in the JWT headers
+  const listResourceTagsResponse = await kms.send(new ListResourceTagsCommand({ KeyId: currentKeyId }))
+  const kid = getTagValueFromTags('jwk_kid', listResourceTagsResponse.Tags ?? [])
+
+  if (kid == null) {
+    return respond('KMS key is not correctly tagged', 500)
+  }
+
+  const iss = process.env.ISSUER
+
+  // JWT Token headers
+  const headers: any = {
+    alg: 'RS256',
+    typ: 'JWT',
+    kid: `${kid}`
+  }
+
+  // prepare token lifetime property values
+  const issuedAtDate = new Date()
+  const expirationDate = new Date(issuedAtDate)
+  const notBeforeDate = new Date(issuedAtDate)
+  expirationDate.setTime(expirationDate.getTime() + 60 * 60 * 1000) // valid for one hour
+  notBeforeDate.setTime(notBeforeDate.getTime() - 5 * 60 * 1000) // 5m before issuedAtDate
+
+  // JWT Token payload
+  const payload: any = {
+    sub: `${identityArn}`, // Set role arn as message for payload
+    aud,
+    iss,
+    iat: Math.floor(issuedAtDate.getTime() / 1000),
+    exp: Math.floor(expirationDate.getTime() / 1000),
+    nbf: Math.floor(notBeforeDate.getTime() / 1000)
+  }
+
+  // Prepare message to be signed by KMS
+  const tokenHeaders = base64url(JSON.stringify(headers))
+  const tokenPayload = base64url(JSON.stringify(payload))
+
+  // Sign message with KMS
+  const signResponse = await kms.send(new SignCommand({
+    KeyId: currentKeyId,
+    Message: Buffer.from(`${tokenHeaders}.${tokenPayload}`),
+    SigningAlgorithm: 'RSASSA_PKCS1_V1_5_SHA_256',
+    MessageType: 'RAW'
+  }))
+  logger.debug(JSON.stringify(signResponse))
+
+  const signature = Buffer
+    .from(signResponse.Signature as Uint8Array)
+    .toString('base64')
+    .replace(/\+/g, '-')
+    .replace(/\//g, '_')
+    .replace(/=/g, '')
+
+  const token = `${tokenHeaders}.${tokenPayload}.${signature}`
+  logger.debug(token)
+
+  return respond(JSON.stringify({
+    token
+  }))
+}
+
+function respond (message: string, statusCode: number = 200) {
+  return {
+    statusCode,
+    body: message
+  }
+}
+
+function getARNFromIdentity (identityArn: string | null) {
+  if (identityArn === undefined || identityArn === null) {
+    return null
+  }
+
+  // Regex for converting arn to base role
+  const captGroups = [
+    'arn:aws:sts:',
+    '(?<regionName>[^:]*)', // group 1
+    ':',
+    '(?<accountId>\\d{12})', // group 2
+    ':assumed-role\\/',
+    '(?<roleName>[A-z0-9\\-]+?)', // group 3
+    '\\/',
+    '(?<user>[^:]*)', // group 4
+    '$'
+  ]
+
+  const regex = new RegExp(captGroups.join(''))
+  const { regionName, accountId, roleName } = regex.exec(identityArn)?.groups ?? {}
+
+  if (regionName === undefined || accountId === undefined || roleName === undefined) {
+    return null
+  }
+
+  // Build base role arn
+  return `arn:aws:iam:${regionName}:${accountId}:role/${roleName}`
+}
+
+function getTagValueFromTags (tagKey: string, tags: Tag[]) {
+  for (const tag of tags) {
+    if (tag.TagKey === tagKey) {
+      return tag.TagValue
+    }
+  }
+
+  return null
+}