initial

Author: Joe Snell

.gitignore

@@ -0,0 +1,2 @@
+.vscode
+.DS_Store

README.md

@@ -0,0 +1,6 @@
+*/__pycache__/*
+*.pyc
+
+requirements.txt
+build
+extension.zip

python/.gitignore

@@ -0,0 +1,12 @@
+[[source]]
+name = "pypi"
+url = "https://pypi.org/simple"
+verify_ssl = true
+
+[dev-packages]
+boto3 = "*"
+
+[packages]
+
+[requires]
+python_version = "3.8"

python/Pipfile

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+set -e
+
+[ -d build ] && rm -rf build
+[ -f extension.zip ] && rm extension.zip
+
+pipenv lock -r > requirements.txt
+pip install -r requirements.txt -t build
+
+chmod +x extensions/parameter-store-extension
+chmod +x parameter-store-extension/extension.py
+
+mkdir -p build/opt/extensions
+
+cp -R extensions build/opt/
+cp -R parameter-store-extension build/opt/
+
+zip -r extension.zip build
+
+aws lambda publish-layer-version \
+    --layer-name "lambda-extensions_parameter-store" \
+    --region "us-east-1" \
+    --zip-file "fileb://extension.zip"

python/Pipfile.lock

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+set -e
+
+bin/build.sh
+
+aws lambda publish-layer-version \
+    --layer-name "lambda-extensions_parameter-store" \
+    --region "us-east-1" \
+    --zip-file "fileb://extension.zip" \
+    --compatible-runtimes python3.8

python/bin/build.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+set -euo pipefail
+
+OWN_FILENAME="$(basename $0)"
+LAMBDA_EXTENSION_NAME="$OWN_FILENAME" # (external) extension name has to match the filename
+
+echo "${LAMBDA_EXTENSION_NAME}  launching extension"
+exec "/opt/${LAMBDA_EXTENSION_NAME}/extension.py"

python/bin/deploy.sh

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+import json
+import os
+import boto3
+import requests
+import signal
+import sys
+from pathlib import Path
+
+
+# global variables
+# extension name has to match the file's parent directory name)
+LAMBDA_EXTENSION_NAME = Path(__file__).parent.name
+
+
+# custom extension code
+def execute_custom_processing(event):
+    # perform custom per-event processing here
+    print(f"[{LAMBDA_EXTENSION_NAME}] Received event: {json.dumps(event)}", flush=True)
+
+
+# boiler plate code
+def handle_signal(signal, frame):
+    # if needed pass this signal down to child processes
+    print(f"[{LAMBDA_EXTENSION_NAME}] Received signal={signal}. Exiting.", flush=True)
+    sys.exit(0)
+
+
+def register_extension():
+    print(f"[{LAMBDA_EXTENSION_NAME}] Registering...", flush=True)
+    headers = {
+        'Lambda-Extension-Name': LAMBDA_EXTENSION_NAME,
+    }
+    payload = {
+        'events': [
+            'INVOKE',
+            'SHUTDOWN'
+        ],
+    }
+    response = requests.post(
+        url=f"http://{os.environ['AWS_LAMBDA_RUNTIME_API']}/2020-01-01/extension/register",
+        json=payload,
+        headers=headers
+    )
+    ext_id = response.headers['Lambda-Extension-Identifier']
+    print(f"[{LAMBDA_EXTENSION_NAME}] Registered with ID: {ext_id}", flush=True)
+
+    return ext_id
+
+
+def process_events(ext_id, client):
+    headers = {
+        'Lambda-Extension-Identifier': ext_id
+    }
+
+    while True:
+        print(f"[{LAMBDA_EXTENSION_NAME}] Waiting for event...", flush=True)
+
+        response = requests.get(
+            url=f"http://{os.environ['AWS_LAMBDA_RUNTIME_API']}/2020-01-01/extension/event/next",
+            headers=headers,
+            timeout=None
+        )
+        event = json.loads(response.text)
+        if event['eventType'] == 'SHUTDOWN':
+            print(f"[{LAMBDA_EXTENSION_NAME}] Received SHUTDOWN event. Exiting.", flush=True)
+            sys.exit(0)
+        else:
+            execute_custom_processing(event)
+
+client = None
+
+def main():
+    global client
+    if client is None:
+        client = boto3.client('ssm')
+
+    # handle signals
+    signal.signal(signal.SIGINT, handle_signal)
+    signal.signal(signal.SIGTERM, handle_signal)
+
+    # execute extensions logic
+    extension_id = register_extension()
+    process_events(extension_id, client)
+
+
+if __name__ == "__main__":
+    main()