(event)
| 18 | lambda_client = boto3.client('lambda') |
| 19 | |
| 20 | def update_function(event): |
| 21 | user_pool_id = event['ResourceProperties']['UserPoolId'] |
| 22 | cognito_region = event['ResourceProperties']['CognitoRegion'] |
| 23 | source_url = event['ResourceProperties'].get('SourceUrl') |
| 24 | edge_function_arn = event['ResourceProperties']['EdgeFunctionArn'] |
| 25 | function_filename = event['ResourceProperties'].get('FunctionFilename', 'index.js') |
| 26 | |
| 27 | logger.info("Downloading well-known jwks.json from Cognito") |
| 28 | jwks_url = f'https://cognito-idp.{cognito_region}.amazonaws.com/{user_pool_id}/.well-known/jwks.json' |
| 29 | with urlopen(jwks_url) as http_response: |
| 30 | jwks = str(http_response.read()) |
| 31 | |
| 32 | jwks = jwks.replace('b\'{', '{') |
| 33 | jwks = jwks.replace('}\'', '}') |
| 34 | logger.debug(json.dumps(jwks, indent = 2, default = str)) |
| 35 | |
| 36 | if not source_url: |
| 37 | logger.info('SourceUrl not specified so determining code location from Lambda for "Templated" alias') |
| 38 | # The "Templated" alias is created when the edge auth function is deployed and represents the original |
| 39 | # version of the function that is templated with replacement variables. |
| 40 | response = lambda_client.get_function( |
| 41 | FunctionName = f'{edge_function_arn}:Templated' |
| 42 | ) |
| 43 | |
| 44 | source_url = response['Code']['Location'] |
| 45 | |
| 46 | logger.info("Building updated function zip archive") |
| 47 | js = None |
| 48 | with urlopen(source_url) as zip_resp: |
| 49 | with zipfile.ZipFile(io.BytesIO(zip_resp.read())) as zin: |
| 50 | with zipfile.ZipFile('/tmp/edge-code.zip', 'w') as zout: |
| 51 | zout.comment = zin.comment |
| 52 | for item in zin.infolist(): |
| 53 | if item.filename == function_filename: |
| 54 | js = io.TextIOWrapper(io.BytesIO(zin.read(item.filename))).read() |
| 55 | else: |
| 56 | zout.writestr(item, zin.read(item.filename)) |
| 57 | |
| 58 | if not js: |
| 59 | raise Exception(f'Function code archive does not contain the file "{function_filename}"') |
| 60 | |
| 61 | js = js.replace('##JWKS##', jwks) |
| 62 | js = js.replace('##USERPOOLID##', user_pool_id) |
| 63 | js = js.replace('##COGNITOREGION##', cognito_region) |
| 64 | |
| 65 | logger.info('Writing updated js file %s to archive', function_filename) |
| 66 | with zipfile.ZipFile('/tmp/edge-code.zip', mode='a', compression=zipfile.ZIP_DEFLATED) as zf: |
| 67 | zf.writestr(function_filename, js) |
| 68 | |
| 69 | # Load file into memory |
| 70 | with open('/tmp/edge-code.zip', 'rb') as file_data: |
| 71 | bytes_content = file_data.read() |
| 72 | |
| 73 | logger.info('Updating lambda function with updated code archive') |
| 74 | response = lambda_client.update_function_code( |
| 75 | FunctionName = edge_function_arn, |
| 76 | ZipFile = bytes_content |
| 77 | ) |
no test coverage detected