Merge pull request #6 from aws-iot-builder-tools/flow-refactor

Flow refactor complete

Author: rpcme

coverage.svg

@@ -19,7 +19,6 @@
 from cryptography.hazmat.backends import default_backend
 from aws_lambda_powertools.utilities.typing import LambdaContext
 from aws_lambda_powertools.utilities.data_classes import SQSEvent
-from aws_utils import get_thing_type_arn, get_thing_group_arn, get_policy_arn
 from cert_utils import get_certificate_fingerprint
 
 logger = logging.getLogger()
@@ -62,6 +61,8 @@ def get_thing(thing_name: str) -> str:
 
 def process_policy(policy_name, certificate_id):
     """Attaches the IoT policy to the certificate"""
+    if policy_name is None:
+        return
     iot_client = boto3client('iot')
     iot_client.attach_policy(policyName=policy_name,
                              target=get_certificate_arn(certificate_id))
@@ -145,15 +146,13 @@ def process_certificate(config, requeue_cb):
                   "Role. Unable to reach IoT Core object.")
         return None
 
-def process_thing_group(thing_group_name, thing_name):
+def process_thing_group(thing_group_arn, thing_arn):
     """Attaches the configured thing group to the iot thing"""
+    if thing_group_arn is None:
+        return
     iot_client = boto3client('iot')
     try:
-        thing_group_arn = get_thing_group_arn(thing_group_name)
-        thing_arn = get_thing(thing_name)
-        iot_client.add_thing_to_thing_group(thingGroupName=thing_group_name,
-                                            thingGroupArn=thing_group_arn,
-                                            thingName=thing_name,
+        iot_client.add_thing_to_thing_group(thingGroupArn=thing_group_arn,
                                             thingArn=thing_arn,
                                             overrideDynamicGroups=False)
     except ClientError as error:
@@ -178,32 +177,18 @@ def get_name_from_certificate(certificate_id):
 
 def process_sqs(config):
     """Main processing function to procedurally run through processing steps."""
-    policy_name = config.get('policy_name')
-    thing_group_name = config.get('thing_group_name')
-    thing_type_name = config.get('thing_type_name')
-
     certificate_id = process_certificate(config, requeue)
-
-    if certificate_id is None:
-        thing_name = config['thing']
-    else:
-        thing_name = get_name_from_certificate(certificate_id)
-
-    process_thing(thing_name, certificate_id, thing_type_name)
-    process_policy(policy_name, certificate_id)
-    process_thing_group(thing_group_name, thing_name)
-
-def lambda_handler(event: SQSEvent, context: LambdaContext) -> dict: # pylint: disable=unused-argument
+    process_thing(config.get('thing'),
+                  certificate_id,
+                  config.get('thing_type_arn'))
+    process_policy(config.get('policy_arn'),
+                   certificate_id)
+    process_thing_group(config.get('thing_group_arn'),
+                        config.get('thing'))
+
+def lambda_handler(event: SQSEvent,
+                   context: LambdaContext) -> dict: # pylint: disable=unused-argument
     """Lambda function main entry point"""
-    if event.get('Records') is None:
-        print("ERROR: Configuration incorrect: no event record on invoke")
-        return {
-            "statusCode": 400,
-            "body": json.dumps({
-                "message": "No SQS event records available for processing."
-            })
-        }
-
     for record in event['Records']:
         if record.get('eventSource') == 'aws:sqs':
             config = json.loads(record["body"])

linting.svg

@@ -1,2 +1,4 @@
-boto3
-cryptography
+botocore==1.38.15
+boto3==1.38.15
+cryptography==44.0.3
+aws_lambda_powertools==3.12.0

src/bulk_importer/main.py

@@ -46,15 +46,11 @@ def s3_object_bytes(bucket_name: str, object_name: str, getvalue: bool=False):
 #                                     object_name=object_name)
 #    return object_stream.getvalue()
 
-def queue_manifest_certificate(identity, certificate, queue_url):
+def send_sqs_message(config, queue_url):
     """Send the thing name and certificate to sqs queue"""
     sqs_client = client("sqs")
-    payload = {
-        'thing': identity,
-        'certificate': certificate
-    }
     sqs_client.send_message( QueueUrl=queue_url,
-                             MessageBody=dumps(payload) )
+                             MessageBody=dumps(config) )
 
 def verify_queue(queue_url: str) -> bool:
     """Verify the queue exists by attempting to fetch its attributes"""
@@ -72,6 +68,8 @@ def verify_queue(queue_url: str) -> bool:
 
 def get_thing_type_arn(type_name: str) -> str:
     """Retrieves the thing type ARN"""
+    if type_name in ("None", ""):
+        return None
     iot_client = client('iot')
     try:
         response = iot_client.describe_thing_type(thingTypeName=type_name)
@@ -83,8 +81,10 @@ def get_thing_type_arn(type_name: str) -> str:
         logger.error("{this} ({thing_type_name}): {error_code} : {error_mesg}")
         raise error
 
-def get_thing_group_arn(thing_group_name):
+def get_thing_group_arn(thing_group_name: str) -> str:
     """Retrieves the thing group ARN"""
+    if thing_group_name in ("None", ""):
+        return None
     iot_client = client('iot')
 
     try:
@@ -97,8 +97,11 @@ def get_thing_group_arn(thing_group_name):
         logger.error("{this} ({thing_group_name}): {error_code} : {error_mesg}")
         raise error
 
-def get_policy_arn(policy_name):
+def get_policy_arn(policy_name: str) -> str:
     """Retrieve the IoT policy ARN"""
+    if policy_name is None:
+        return None
+
     iot_client = client('iot')
     try:
         response = iot_client.get_policy(policyName=policy_name)

src/bulk_importer/requirements.txt

@@ -7,43 +7,87 @@
 import os
 import json
 import boto3
-from aws_lambda_powertools.utilities.typing import LambdaContext
-from aws_lambda_powertools.utilities.data_classes import SQSEvent
-
-def process(payload):
-    """Annotate payload with environment-passed variants, later this function
-       will evolve to allow importing types, groups, and policies"""
+import logging
 
-    queue_url = os.environ.get('QUEUE_TARGET')
+from aws_lambda_powertools.utilities.typing import LambdaContext
+from aws_lambda_powertools.utilities.data_classes import S3Event
+from aws_utils import get_policy_arn, get_thing_group_arn, get_thing_type_arn
 
-    # Policy is required.
-    payload['policy_name'] = os.environ.get('POLICY_NAME')
+logger = logging.getLogger()
+logger.setLevel("INFO")
 
-    # Thing group is desired, but optional.
-    # The reason why 'None' has to be set is an environment variable
-    # on a Lambda function cannot be set to empty
-    if os.environ.get('THING_GROUP_NAME') == "None":
-        payload['thing_group_name'] = None
-    else:
-        payload['thing_group_name'] = os.environ.get('THING_GROUP_NAME')
+ESPRESSIF_BUCKET_PREFIX = "thingpress-espressif-"
+INFINEON_BUCKET_PREFIX = "thingpress-infineon-"
+MICROCHIP_BUCKET_PREFIX = "thingpress-microchip-"
 
-    # Thing group is desired, but optional.
-    if os.environ.get('THING_TYPE_NAME') == "None":
-        payload['thing_type_name'] = None
-    else:
-        payload['thing_type_name'] = os.environ.get('THING_TYPE_NAME')
+def process(payload: hash, queue_url: str) -> hash:
+    """Annotate payload with environment-passed variants, later this function
+       will evolve to allow importing types, groups, and policies"""
 
     # Pass on to the queue for target processing.
     client = boto3.client("sqs")
+    logger.info("Sending payload to queue {queue_url}")
     client.send_message( QueueUrl=queue_url,
                          MessageBody=json.dumps(payload))
     return payload
 
-def lambda_handler(event: SQSEvent, context: LambdaContext) -> dict: # pylint: disable=unused-argument
-    """Lambda function main entry point"""
+def get_provider_queue(bucket_name: str) -> str:
+    """
+    Returns the queue related to the prefix of a given bucket
+    The cfn stack prescribes the environment variable value.
+    See the cfn template for more detail.
+    """
+    if bucket_name.startswith(ESPRESSIF_BUCKET_PREFIX):
+        return os.environ.get('QUEUE_TARGET_ESPRESSIF')
+    if bucket_name.startswith(INFINEON_BUCKET_PREFIX):
+        return os.environ.get('QUEUE_TARGET_INFINEON')
+    if bucket_name.startswith(MICROCHIP_BUCKET_PREFIX):
+        return os.environ.get('QUEUE_TARGET_MICROCHIP')
+    return None
+
+def lambda_handler(event: S3Event,
+                   context: LambdaContext) -> dict: # pylint: disable=unused-argument
+    """
+    Lambda function main entry point. Verifies the S3 object can be read and resolves
+    inputs prior to forwarding to vendor handler queue.
+
+    This lambda function expects invocation by S3 event. There should be only one
+    event, but is processed as if multiple events were found at once.
+    
+    Expects the following environment variables to be set:
+    QUEUE_TARGET_ESPRESSIF
+    QUEUE_TARGET_INFINEON
+    QUEUE_TARGET_MICROCHIP
+    
+    Expects at least one of the following environment variables to be set:
+    POLICY_NAME
+    THING_GROUP_NAME
+
+    May have the following environment variables set:
+    THING_TYPE_NAME
+    """
     # Get the payload coming in and process it.  There might be more than one.
-    result = []
-    for record in event['Records']:
-        r = process(json.loads(record["body"]))
-        result.append(r)
-    return result
+    v_thing_group = get_thing_group_arn(os.environ.get('THING_GROUP_NAME'))
+    v_thing_type = get_thing_type_arn(os.environ.get('THING_TYPE_NAME'))
+    v_policy = get_policy_arn(os.environ.get('POLICY_NAME'))
+
+    s3_event = S3Event(event)
+    queue_url = get_provider_queue(s3_event.bucket_name)
+    if queue_url is None:
+        logger.error("Queue URL could not be resolved. Exiting.")
+        return None
+
+    data = {
+        'policy_arn': os.environ.get('POLICY_NAME'),
+        'thing_group_arn': v_thing_group,
+        'thing_type_arn': v_thing_type,
+        'bucket': s3_event.bucket_name
+    }
+
+    for record in s3_event.records:
+        # TODO: verify s3 object, for now assume it is reachable
+        # v_object = verify_s3_object(bucket, record.s3.get_object.key)
+        data['key'] = record.s3.get_object.key
+        logger.info("Processing data for object {record.s3.get_object.key}")
+        process(data, queue_url)
+    return event

src/layer_utils/aws_utils.py

@@ -8,25 +8,38 @@
 import os
 import io
 import csv
+import json
+import base64
+import logging
 from aws_lambda_powertools.utilities.typing import LambdaContext
-from aws_lambda_powertools.utilities.data_classes import S3Event
-from aws_utils import s3_object_bytes, queue_manifest_certificate
+from aws_lambda_powertools.utilities.data_classes import SQSEvent
+from aws_utils import s3_object_bytes, send_sqs_message
 
-def invoke_export(bucket_name: str, object_name: str, queue_url: str):
+logger = logging.getLogger()
+logger.setLevel("INFO")
+
+def invoke_export(config: hash, queue_url: str):
     """Evaluate CSV based Espressif manifest"""
-    manifest_bytes = s3_object_bytes(bucket_name, object_name, getvalue=True)
+    manifest_bytes = s3_object_bytes(config['bucket'],
+                                     config['key'],
+                                     getvalue=True)
     reader_list = csv.DictReader(io.StringIO(manifest_bytes.decode()))
 
     for row in reader_list:
-        queue_manifest_certificate(row['MAC'], row['cert'], queue_url)
+        config['thing'] = row['MAC']
+        config['certificate'] = str(base64.b64encode(row['cert'].encode('ascii')))
+        send_sqs_message(config, queue_url)
 
 def lambda_handler(event: dict, context: LambdaContext) -> dict: # pylint: disable=unused-argument
     """Lambda function main entry point"""
-    s3_event = S3Event(event)
+    sqs_event = SQSEvent(event)
     queue_url = os.environ['QUEUE_TARGET']
+    if event.get('Records') is None:
+        #TODO throw an exception here
+        return None
+    for record in event['Records']:
+        if record.get('eventSource') == 'aws:sqs':
+            config = json.loads(record["body"])
+            invoke_export(config, queue_url)
 
-    bucket = s3_event.bucket_name
-    for record in s3_event.records:
-        manifest = record.s3.get_object.key
-        invoke_export(bucket, manifest, queue_url)
     return event

src/product_provider/main.py

@@ -6,10 +6,11 @@
 the import processing pipeline
 """
 import os
+import json
 import logging
 from aws_lambda_powertools.utilities.typing import LambdaContext
-from aws_lambda_powertools.utilities.data_classes import S3Event
-from aws_utils import s3_object_bytes, verify_queue
+from aws_lambda_powertools.utilities.data_classes import SQSEvent
+from aws_utils import verify_queue
 from .manifest_handler import invoke_export, verify_certtype
 
 logger = logging.getLogger()
@@ -28,10 +29,14 @@ def lambda_handler(event: dict, context: LambdaContext) -> dict: # pylint: disab
         logger.error("Certificate type not valid. Must be E0E0, E0E1, or E0E2.")
         return None
 
-    s3_event = S3Event(event)
-    bucket = s3_event.bucket_name
-    for record in s3_event.records:
-        manifest = record.s3.get_object.key
-        manifest_content = s3_object_bytes(bucket, manifest, getvalue=True)
-        invoke_export(manifest_content, queue_url, cert_type)
+    sqs_event = SQSEvent(event)
+    queue_url = os.environ['QUEUE_TARGET']
+    if event.get('Records') is None:
+        #TODO throw an exception here
+        return None
+    for record in event['Records']:
+        if record.get('eventSource') == 'aws:sqs':
+            config = json.loads(record["body"])
+            invoke_export(config, queue_url, cert_type)
+
     return event