├── .gitignore ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── NOTICE ├── README.md ├── cloudformation ├── function.yaml └── sam-bootstrap.yaml ├── images ├── ref-arch.png └── workflow.png ├── pytest.ini ├── requirements_validate_scan.txt ├── scripts ├── lint.sh ├── sam.sh ├── scan.sh └── test.sh ├── src ├── __init__.py ├── helper.py ├── main.py └── requirements.txt ├── test ├── __init__.py ├── local_send_message_test.py ├── local_send_message_test.yaml ├── sqs_messages │ ├── test_minimum_requirements.json │ ├── test_network_override.json │ ├── test_no_container_name.json │ ├── test_no_network.json │ └── test_no_security_group.json └── test_ecs_run_task.py └── tox.ini /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | **/generated-sam-template.yaml 3 | **/.aws-sam 4 | **/test-event.json 5 | **/.python-version 6 | __pycache__/ 7 | scripts/logs 8 | .coverage 9 | .venv 10 | htmlcov/ 11 | coverage.xml 12 | junitxml.xml 13 | *.zip 14 | /logs 15 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | ## Code of Conduct 2 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 3 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 4 | opensource-codeofconduct@amazon.com with any additional questions or comments. 5 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing Guidelines 2 | 3 | Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional 4 | documentation, we greatly value feedback and contributions from our community. 5 | 6 | Please read through this document before submitting any issues or pull requests to ensure we have all the necessary 7 | information to effectively respond to your bug report or contribution. 8 | 9 | 10 | ## Reporting Bugs/Feature Requests 11 | 12 | We welcome you to use the GitHub issue tracker to report bugs or suggest features. 13 | 14 | When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already 15 | reported the issue. Please try to include as much information as you can. Details like these are incredibly useful: 16 | 17 | * A reproducible test case or series of steps 18 | * The version of our code being used 19 | * Any modifications you've made relevant to the bug 20 | * Anything unusual about your environment or deployment 21 | 22 | 23 | ## Contributing via Pull Requests 24 | Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that: 25 | 26 | 1. You are working against the latest source on the *main* branch. 27 | 2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already. 28 | 3. You open an issue to discuss any significant work - we would hate for your time to be wasted. 29 | 30 | To send us a pull request, please: 31 | 32 | 1. Fork the repository. 33 | 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change. 34 | 3. Ensure local tests pass. 35 | 4. Commit to your fork using clear commit messages. 36 | 5. Send us a pull request, answering any default questions in the pull request interface. 37 | 6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation. 38 | 39 | GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and 40 | [creating a pull request](https://help.github.com/articles/creating-a-pull-request/). 41 | 42 | 43 | ## Finding contributions to work on 44 | Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start. 45 | 46 | 47 | ## Code of Conduct 48 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 49 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 50 | opensource-codeofconduct@amazon.com with any additional questions or comments. 51 | 52 | 53 | ## Security issue notifications 54 | If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue. 55 | 56 | 57 | ## Licensing 58 | 59 | See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution. 60 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | -------------------------------------------------------------------------------- /NOTICE: -------------------------------------------------------------------------------- 1 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # aws-lambda-amazon-ecs-run-task 2 | 3 | ### Description: 4 | This function will allow you to run an AWS ECS Task from an SQS Queue. 5 | 6 | ### Folder Structure 7 | | Folder/File | Description | 8 | | :-------------------------| :-------------------------------------------------------------------------------------------------------------------| 9 | | cloudformation/sam-bootstrap.yaml | AWS Cloudformation template that will create the required AWS Resources for the solution to work properly. It will create an IAM Role, KMS Key/Alias and S3 Bucket. All of these AWS Resources are required for an AWS Serverless Application Model (SAM) deployment to successful.| 10 | | cloudformation/function.yaml | SAM template that will deploy the AWS Lambda Function along with all dependant infrastructure. | 11 | | src | Source code for AWS Lambda Functions. | 12 | | test | Test code for AWS Lambda Functions. | 13 | | scripts | Directory that has the scripts that will be ran from a CodeBuild BuildSpec file | 14 | | scripts/main.sh | An orchestration script that will run the all other linting/scanning scripts before building/deploying the SAM Function(s). | 15 | | scripts/pylint.sh | Shell script that will run the ```pylint``` command against all python files. | 16 | | scripts/pyscan.sh | Runs Bandit (python lib) against all python code within the repository to identify any security vulnerabilities. | 17 | | scripts/sam.sh | Runs a number of SAM commands to package / build / deploy the SAM Function to a specified account. | 18 | | scripts/test.sh | Shell script that will run the ```tox``` command to build a virtual environment and the ```pytest``` command to run any unit tests found in the repository. | 19 | | pytest.ini | ini files are the configuration files of the tox project, and can also be used to hold pytest configuration if they have a [pytest] section. | 20 | | test_requirements.txt | Pip requirements file for test environment. | 21 | | tox.ini | Configured file for Tox. Tox is a command-line driven automated testing tool for Python, based on the use of virtualenv. | 22 | 23 | ## Pre-requisite Steps: 24 | - Install the Serverless Application Model CLI (SAM) [Link to AWS Doc](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-sam-cli-install.html) 25 | - Since this solution builds the SAM function inside Lambda-like container, Docker must be installed and running on your workstation. 26 | 27 | ## How to deploy 28 | ### AWS Management or Shared Services Production Account 29 | - Run the cloudformation/sam-bootstrap.yaml into the AWS Account you chose. 30 | 31 | ```bash 32 | aws cloudformation create-stack --stack-name SAM-Bootstrap --template-body file://cloudformation/sam-bootstrap.yaml 33 | ``` 34 | 35 | - Deploy Serverless Application Model function. 36 | ```bash 37 | bash scripts/sam.sh 38 | ```` 39 | 40 | ## How does the Lambda ECS Run Task work 41 | 42 | ### Reference Architecture 43 | ![alt text](images/ref-arch.png) 44 | 45 | ### Workflow 46 | ![alt text](images/workflow.png) 47 | 48 | ### Input Parameters / Configuration 49 | message_body ECS 50 | 51 | * **cluster** (*string*) -- 52 | 53 | The short name or full Amazon Resource Name (ARN) of the cluster on which to run your task. If you do not specify a cluster, the default cluster is assumed. 54 | 55 | * **task_def** (*string*) -- 56 | 57 | The family and revision (family:revision ) or full ARN of the task definition to run. If a revision is not specified, the latest ACTIVE revision is used. 58 | 59 | * **container_name** (*string*) -- 60 | 61 | The name of the container that receives the override. This parameter is required if any override is specified. 62 | 63 | * **environment** (*list*) -- 64 | 65 | The environment variables to send to the container. You can add new environment variables, which are added to the container at launch, or you can override the existing environment variables from the Docker image or the task definition. You must also specify a container name. 66 | 67 | (dict) -- 68 | A key-value pair object. 69 | 70 | * **Name** (*string*) -- 71 | 72 | The name of the key-value pair. For environment variables, this is the name of the environment variable. 73 | 74 | * **Value** (*string*) -- 75 | 76 | The value of the key-value pair. For environment variables, this is the value of the environment variable. 77 | 78 | * **command** (*list*) -- 79 | 80 | The command to send to the container that overrides the default command from the Docker image or the task definition. You must also specify a container name. 81 | 82 | * **subnets** (*list*) -- 83 | 84 | The IDs of the subnets associated with the task or service. There is a limit of 16 subnets that can be specified per AwsVpcConfiguration 85 | 86 | * **security_groups** (*list*) -- 87 | 88 | The IDs of the security groups associated with the task or service. If you do not specify a security group, the default security group for the VPC is used. There is a limit of 5 security groups that can be specified per AwsVpcConfiguration 89 | 90 | * **cpu** (*string*) -- 91 | 92 | The cpu override for the task. 93 | 94 | * **memory** (*string*) -- 95 | 96 | The memory override for the task. 97 | 98 | #### SQS Message Example Code [JSON]: 99 | To easily present the ECS Task data to the Lambda Function the message must into a dictionary using the ECS key. 100 | 101 | ```json 102 | message_body['ECS'] = { 103 | "cluster": "container-test01", 104 | "taskDefinition": "container-test01:6", 105 | "service": "container-test01", 106 | "launchType": "FARGATE", 107 | "platformVersion": "LATEST", 108 | "propagateTags": "TASK_DEFINITION", 109 | "overrides": { 110 | "containerOverrides": [ 111 | { 112 | "name": "container-test01-development", 113 | "command": ["sh", "-c", "sleep 120 && df -h"], 114 | "environment": [ 115 | {"name": "s3_image_id", "value": "00000.tiff"}, 116 | {"name": "s3_bucket", "value": "testapp22-111122223333-us-east-1"}, 117 | {"name": "sleep_time", "value": "1"} 118 | ] 119 | } 120 | ], 121 | "cpu": "1024", 122 | "memory": "2048" 123 | }, 124 | "networkConfiguration": { 125 | "awsvpcConfiguration": { 126 | "subnets": ["subnet-1234567890abcdef0", "subnet-abcdef01234567890"], 127 | "securityGroups": ["sg-021345abcdef6789"], 128 | "assignPublicIp": "ENABLED" 129 | } 130 | }, 131 | "referenceId": ""} 132 | } 133 | ``` 134 | 135 | #### SQS Message Example Code [JSON]: 136 | Before sending the message to the queue ensure the message body is converted to a string. 137 | 138 | ```python 139 | ... 140 | 141 | queue_url = f"https://sqs.us-east-1.amazonaws.com/111122223333/EcsRunTask" 142 | response = sqs.send_message( 143 | QueueUrl=queue_url, 144 | MessageBody=json.dumps(message_body) 145 | ) 146 | 147 | ... 148 | 149 | ``` 150 | 151 | #### Running Local Send Message Test 152 | - Update test/local_send_message_test.yaml with environment information (i.e. cluster name / service / etc) 153 | - Run local_send_message_test.py script 154 | ```bash 155 | python3 test/local_send_message_test.py 156 | ``` 157 | -------------------------------------------------------------------------------- /cloudformation/function.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | Transform: "AWS::Serverless-2016-10-31" 5 | Description: This function will trigger an ECS Task from an SQS Message. 6 | 7 | Parameters: 8 | pApplicationName: 9 | Type: String 10 | Default: "EcsRunTask" 11 | pSecurityGroup: 12 | Type: String 13 | Default: "" 14 | 15 | Resources: 16 | # ------------------- 17 | # IAM Role / Policies 18 | # ------------------- 19 | rLambdaFunctionRole: 20 | Type: AWS::IAM::Role 21 | Properties: 22 | RoleName: !Sub ${pApplicationName}-Lambda-Role 23 | AssumeRolePolicyDocument: 24 | Version: '2012-10-17' 25 | Statement: 26 | - Effect: Allow 27 | Principal: 28 | Service: 29 | - lambda.amazonaws.com 30 | Action: 31 | - sts:AssumeRole 32 | Path: "/" 33 | ManagedPolicyArns: 34 | - arn:aws:iam::aws:policy/service-role/AWSLambdaSQSQueueExecutionRole 35 | - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole 36 | - arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy 37 | Policies: 38 | - PolicyName: !Sub ${pApplicationName}-Lambda-Policy 39 | PolicyDocument: 40 | Version: 2012-10-17 41 | Statement: 42 | - Action: 43 | - iam:PassRole 44 | Resource: "*" 45 | Effect: Allow 46 | 47 | # ----------------- 48 | # Lambda Function 49 | # ----------------- 50 | rLambdaFunction: 51 | Type: AWS::Serverless::Function 52 | Properties: 53 | Handler: main.lambda_handler 54 | Runtime: python3.7 55 | FunctionName: !Ref pApplicationName 56 | Description: This function will trigger an ECS Task from an SQS Message. 57 | Timeout: 330 58 | CodeUri: ../src 59 | Role: !GetAtt rLambdaFunctionRole.Arn 60 | Environment: 61 | Variables: 62 | FAILURE_SNS_TOPIC_ARN: !Ref rFailureSNSTopic 63 | QUEUE_URL: !Ref rLambdaFunctionQueue 64 | SQS_RETRY_LIMIT: 6 65 | BOTO_MAX_ATTEMPTS: 10 66 | ECS_SECURITY_GROUP: !Ref pSecurityGroup 67 | Events: 68 | MySQSEvent: 69 | Type: SQS 70 | Properties: 71 | Queue: !GetAtt rLambdaFunctionQueue.Arn 72 | BatchSize: 1 73 | 74 | # ------- 75 | # KMS 76 | # ------- 77 | rLambdaFunctionKmsKey: 78 | Type: AWS::KMS::Key 79 | Properties: 80 | Description: Used for the ecs-run-task sqs. 81 | EnableKeyRotation: true 82 | KeyPolicy: 83 | Version: '2012-10-17' 84 | Id: LambdaFunctionKmsKey 85 | Statement: 86 | - Sid: Allow access for Key Administrators 87 | Effect: Allow 88 | Principal: 89 | AWS: 90 | - !Sub arn:aws:iam::${AWS::AccountId}:root 91 | - !GetAtt rLambdaFunctionRole.Arn 92 | Action: 93 | - kms:* 94 | Resource: "*" 95 | 96 | # ----------------- 97 | # CloudWatch Logs 98 | # ----------------- 99 | rLambdaFunctionLogs: 100 | Type: AWS::Logs::LogGroup 101 | Properties: 102 | LogGroupName: !Sub "/aws/lambda/${rLambdaFunction}" 103 | RetentionInDays: 7 104 | 105 | # ------ 106 | # SQS 107 | # ------ 108 | rLambdaFunctionQueue: 109 | Type: AWS::SQS::Queue 110 | Properties: 111 | QueueName: !Sub ${pApplicationName}-Queue 112 | KmsMasterKeyId: !Ref rLambdaFunctionKmsKey 113 | VisibilityTimeout: 330 114 | RedrivePolicy: 115 | deadLetterTargetArn: !GetAtt rLambdaFunctionQueueDlq.Arn 116 | maxReceiveCount: 5 117 | 118 | rLambdaFunctionQueueDlq: 119 | Type: AWS::SQS::Queue 120 | Properties: 121 | QueueName: !Sub ${pApplicationName}-DLQ 122 | 123 | # ------------------ 124 | # SNS for Failures 125 | # ------------------ 126 | rFailureSNSTopic: 127 | Type: AWS::SNS::Topic 128 | Properties: 129 | DisplayName: !Sub ${pApplicationName}-Failure 130 | TopicName: !Sub ${pApplicationName}-Failure 131 | KmsMasterKeyId: !Ref rLambdaFunctionKmsKey 132 | 133 | # -------------------- 134 | # CloudWatch Rule 135 | # -------------------- 136 | rFailedContainerExecution: 137 | Type: AWS::Events::Rule 138 | Properties: 139 | Name: "FailedContainerExecution" 140 | Description: "Re-Run the Lambda ECS Run Task function if hit network interface timeout" 141 | EventPattern: 142 | source: 143 | - "aws.ecs" 144 | detail-type: 145 | - "ECS Task State Change" 146 | detail: 147 | lastStatus: 148 | - STOPPED 149 | stoppedReason: 150 | - "Timeout waiting for network interface provisioning to complete." 151 | State: "ENABLED" 152 | Targets: 153 | - Arn: !GetAtt rLambdaFunction.Arn 154 | Id: "FailedContainerExecution_V1" 155 | 156 | Outputs: 157 | oFailureSNSTopic: 158 | Value: !Ref rFailureSNSTopic -------------------------------------------------------------------------------- /cloudformation/sam-bootstrap.yaml: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | AWSTemplateFormatVersion: "2010-09-09" 5 | Description: Deploys S3 Bucket and KMS Key for the use of the SAM deployment 6 | 7 | Resources: 8 | # ----- 9 | # KMS 10 | # ----- 11 | rSAMKmsKey: 12 | Type: AWS::KMS::Key 13 | Properties: 14 | Description: Used by the SAM S3 Bucket. 15 | EnableKeyRotation: true 16 | KeyPolicy: 17 | Version: '2012-10-17' 18 | Id: !Ref AWS::StackName 19 | Statement: 20 | - Sid: Allow access for Key Administrators 21 | Effect: Allow 22 | Principal: 23 | AWS: 24 | - !Sub arn:aws:iam::${AWS::AccountId}:root 25 | Action: 26 | - kms:* 27 | Resource: "*" 28 | 29 | rSAMKmsKeyAlias: 30 | Type: AWS::KMS::Alias 31 | Properties: 32 | AliasName: alias/Serverless-Application 33 | TargetKeyId: !Ref rSAMKmsKey 34 | 35 | # ----- 36 | # S3 37 | # ----- 38 | rSAMBucket: 39 | Type: AWS::S3::Bucket 40 | Properties: 41 | BucketName: !Sub sam-${AWS::AccountId}-${AWS::Region} 42 | AccessControl: BucketOwnerFullControl 43 | PublicAccessBlockConfiguration: 44 | BlockPublicAcls: true 45 | BlockPublicPolicy: true 46 | IgnorePublicAcls: true 47 | RestrictPublicBuckets: true 48 | BucketEncryption: 49 | ServerSideEncryptionConfiguration: 50 | - ServerSideEncryptionByDefault: 51 | KMSMasterKeyID: !Ref rSAMKmsKey 52 | SSEAlgorithm: aws:kms 53 | VersioningConfiguration: 54 | Status: Enabled 55 | 56 | rSAMBucketPolicy: 57 | Type: AWS::S3::BucketPolicy 58 | Properties: 59 | Bucket: !Ref rSAMBucket 60 | PolicyDocument: 61 | Statement: 62 | - Sid: DenyUnEncryptedObjectUploads 63 | Effect: Deny 64 | Principal: '*' 65 | Action: s3:PutObject 66 | Resource: 67 | - !Sub ${rSAMBucket.Arn}/* 68 | Condition: 69 | StringNotEquals: 70 | s3:x-amz-server-side-encryption: aws:kms 71 | - Sid: DenyInsecureConnections 72 | Effect: Deny 73 | Principal: '*' 74 | Action: s3:* 75 | Resource: 76 | - !Sub ${rSAMBucket.Arn}/* 77 | - !Sub ${rSAMBucket.Arn} 78 | Condition: 79 | Bool: 80 | aws:SecureTransport: 'false' 81 | - Sid: DisableSigV2 82 | Effect: Deny 83 | Principal: '*' 84 | Action: s3:* 85 | Resource: 86 | - !Sub ${rSAMBucket.Arn}/* 87 | - !Sub ${rSAMBucket.Arn} 88 | Condition: 89 | StringEquals: 90 | s3:signatureversion: AWS -------------------------------------------------------------------------------- /images/ref-arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-lambda-ecs-run-task/57cbb2c1ccc409d36aa6da1aff3b540ac4bbe237/images/ref-arch.png -------------------------------------------------------------------------------- /images/workflow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-lambda-ecs-run-task/57cbb2c1ccc409d36aa6da1aff3b540ac4bbe237/images/workflow.png -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | norecursedirs = .git .idea .venv .vscode iac 3 | python_files = test_*.py Test*.py 4 | junit_family=xunit1 5 | -------------------------------------------------------------------------------- /requirements_validate_scan.txt: -------------------------------------------------------------------------------- 1 | # colorama>=0.4.4 2 | # six~=1.15 3 | awscli==1.19.25 4 | # aws-sam-cli==1.20.0 5 | bandit-aws==0.0.20200316120836 6 | bandit==1.6.2 7 | pytest-cov==2.8.1 8 | pytest==5.3.5 9 | pylint==2.8.2 10 | tox==3.23.0 11 | tox-pyenv==1.1.0 12 | safety==1.10.3 13 | cfn-lint==0.49.0 -------------------------------------------------------------------------------- /scripts/lint.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 4 | # SPDX-License-Identifier: Apache-2.0 5 | 6 | OUTPUT_FILE="logs/lint-output.txt" 7 | mkdir "logs" 8 | echo "Output File: "${OUTPUT_FILE}"" 9 | echo -n "" > "${OUTPUT_FILE}" 10 | 11 | echo "#---------------------------------------------------#" | tee -a "${OUTPUT_FILE}" 12 | echo "# Linting Python Files " | tee -a "${OUTPUT_FILE}" 13 | echo "#---------------------------------------------------#" | tee -a "${OUTPUT_FILE}" 14 | find "." -name "*.py" | \ 15 | grep -Ev ".venv|.pytest_cach|.tox|botocore|boto3|.aws" | \ 16 | xargs pylint --rcfile .pylintrc | tee -a "${OUTPUT_FILE}" 17 | 18 | 19 | echo "#---------------------------------------------------#" | tee -a "${OUTPUT_FILE}" 20 | echo "# Linting Cfn Files " | tee -a "${OUTPUT_FILE}" 21 | echo "#---------------------------------------------------#" | tee -a "${OUTPUT_FILE}" 22 | # https://github.com/aws-cloudformation/cfn-python-lint/issues/1265 23 | IGNORED_FILES=( 24 | './iac/CloudFormation/CodePipeline.yaml' 25 | ) 26 | 27 | ALL_CFN_TEMPLATES=$(grep -r '^AWSTemplateFormatVersion' . | cut -d: -f1) 28 | 29 | for TEMPLATE in ${ALL_CFN_TEMPLATES}; do 30 | if [[ "${TEMPLATE}" == "${IGNORED_FILES[0]}" ]] || [[ "${TEMPLATE}" == "${IGNORED_FILES[1]}" ]]; then 31 | echo "Template Ignored: $TEMPLATE" | tee -a "${OUTPUT_FILE}" 32 | continue 33 | fi 34 | echo "Linting CloudFormation Template - ${TEMPLATE}" | tee -a "${OUTPUT_FILE}" 35 | rm -f /tmp/cfn-lint-output.txt 36 | cfn-lint -t "${TEMPLATE}" | tee -a "${OUTPUT_FILE}" 37 | done -------------------------------------------------------------------------------- /scripts/sam.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -e 4 | 5 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 6 | # SPDX-License-Identifier: Apache-2.0 7 | 8 | echo "#--------------------------------------------------------#" 9 | echo "# Building SAM Packages for ${BASE} " 10 | echo "#--------------------------------------------------------#" 11 | 12 | region=$(aws configure get region) || region="us-east-1" 13 | BUCKET=$(aws s3 ls |awk '{print $3}' |grep -E "^sam-[0-9]{12}-${region}" ) 14 | 15 | KMS=$(aws s3api get-bucket-encryption \ 16 | --bucket "${BUCKET}" \ 17 | --region "${region}" \ 18 | --query 'ServerSideEncryptionConfiguration.Rules[*].ApplyServerSideEncryptionByDefault.KMSMasterKeyID' \ 19 | --output text 20 | ) 21 | 22 | echo "Deploying Serverless Application Function" 23 | 24 | sam build -t cloudformation/function.yaml --use-container --region "${region}" 25 | 26 | sam package \ 27 | --template-file .aws-sam/build/template.yaml \ 28 | --s3-bucket "${BUCKET}" \ 29 | --s3-prefix "SAM" \ 30 | --kms-key-id "${KMS}" \ 31 | --region "${region}" \ 32 | --output-template-file cloudformation/generated-sam-template.yaml 33 | 34 | sam deploy \ 35 | --stack-name LambdaEcsRunTask \ 36 | --template-file cloudformation/generated-sam-template.yaml \ 37 | --capabilities CAPABILITY_NAMED_IAM 38 | -------------------------------------------------------------------------------- /scripts/scan.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 4 | # SPDX-License-Identifier: Apache-2.0 5 | 6 | OUTPUT_FILE="logs/scan-output.txt" 7 | echo -n "" > "${OUTPUT_FILE}" 8 | 9 | echo "#-----------------------------------------------#" | tee -a "${OUTPUT_FILE}" 10 | echo "# Scanning Python code for Vulnerabilities " | tee -a "${OUTPUT_FILE}" 11 | echo "#-----------------------------------------------#" | tee -a "${OUTPUT_FILE}" 12 | echo "Output File: "${OUTPUT_FILE}"" 13 | bandit --recursive ../src | tee -a "${OUTPUT_FILE}" 14 | 15 | echo "#-----------------------------------------------#" | tee -a "${OUTPUT_FILE}" 16 | echo "# Scanning Python depend for Vulnerabilities " | tee -a "${OUTPUT_FILE}" 17 | echo "#-----------------------------------------------#" | tee -a "${OUTPUT_FILE}" 18 | safety check | tee -a "${OUTPUT_FILE}" -------------------------------------------------------------------------------- /scripts/test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 4 | # SPDX-License-Identifier: Apache-2.0 5 | 6 | set -eo pipefail 7 | 8 | echo "#---------------------------------------------------#" 9 | echo "# Running Tests " 10 | echo "#---------------------------------------------------#" 11 | 12 | for test in $(find . -name tox.ini); do 13 | #pyenv local 3.6.9 3.7.4 3.8.1 # Can do multiple versions like this but need to make sure they are installed 14 | # pyenv local 3.7.4 15 | tox -c "${test}" 16 | done -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/aws-lambda-ecs-run-task/57cbb2c1ccc409d36aa6da1aff3b540ac4bbe237/src/__init__.py -------------------------------------------------------------------------------- /src/helper.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | import logging 5 | import time 6 | import re 7 | import os 8 | import json 9 | from functools import wraps 10 | from random import randint 11 | import boto3 12 | 13 | logging.basicConfig() 14 | logger = logging.getLogger() 15 | logging.getLogger("botocore").setLevel(logging.ERROR) 16 | logger.setLevel(logging.INFO) 17 | 18 | 19 | def send_failure_email(error_message, add_message=None): 20 | """This will send and email to an SNS topic about a Lambda Function Failure 21 | 22 | Args: 23 | error_message (str): Error message the was produced by the Lambda Function 24 | add_message (str) [optional]: Any additional data that would be added to the Email 25 | 26 | Returns: 27 | N/A 28 | """ 29 | logger.info("Sending error message to SNS Topic") 30 | message = f'''\ 31 | Hello, 32 | 33 | Your Function failed with the following message. 34 | 35 | Error Message: {error_message} 36 | 37 | Additional Data: {add_message} 38 | 39 | To help further debug your issue... 40 | AWS_LAMBDA_LOG_GROUP_NAME: {os.environ['AWS_LAMBDA_LOG_GROUP_NAME']} 41 | AWS_LAMBDA_LOG_STREAM_NAME: {os.environ['AWS_LAMBDA_LOG_STREAM_NAME']} 42 | 43 | Sincerely, 44 | Your friendly neighborhood cloud architect :) 45 | 46 | ''' 47 | 48 | client = boto3.client('sns') 49 | try: 50 | response = client.publish( 51 | TopicArn=os.environ['FAILURE_SNS_TOPIC_ARN'], 52 | Message=message, 53 | Subject=f"Execution of Lambda Function:{os.environ['AWS_LAMBDA_FUNCTION_NAME']} failed." 54 | ) 55 | logger.info(response) 56 | 57 | except Exception as err: 58 | logger.error(err, exc_info=True) 59 | 60 | 61 | def retry(max_attempts: int = 5, delay: int = 3, error_code=None, error_message=None): 62 | """Used as a decorator to retry any definition based on error code or error message, typically used due 63 | to lower then normal API limits. 64 | 65 | Args: 66 | max_attempts (int): Max number of retries 67 | delay (int): Duration to wait before another retry 68 | error_code (list) [optional]: Error code to search for in error 69 | error_message (str) [optional]: Error message to search for to retry 70 | 71 | Returns: 72 | :obj:`json`: Returns json object (dict) of the user parameters 73 | """ 74 | 75 | if not error_code: 76 | error_code = ['TooManyRequestsException'] 77 | 78 | def retry_decorator(function): 79 | @wraps(function) 80 | def wrapper(*args, **kwargs): 81 | last_exception = "default exception" 82 | m_attempts = max_attempts # create new memory allocation 83 | while m_attempts > 1: 84 | try: 85 | logger.debug("**** [retry] args:%s", args) 86 | logger.debug("**** [retry] kwargs:%s", kwargs) 87 | return function(*args, **kwargs) 88 | 89 | except Exception as err: 90 | if hasattr(err, 'response'): 91 | logger.debug("**** [retry] ErrorCode:%s", err.response['Error']['Code']) 92 | 93 | logger.warning(err) 94 | if error_message and re.search(error_message, str(err)): 95 | logger.warning( 96 | "Definition failed:%s with (%s) message, trying again " 97 | "in %s seconds...", function.__name__, error_message, delay 98 | ) 99 | time.sleep(delay) 100 | m_attempts -= 1 101 | last_exception = err 102 | 103 | elif hasattr(err, 'response'): 104 | if err.response['Error']['Code'] in error_code: 105 | logger.warning( 106 | "Definition failed:%s with (%s) " 107 | "error code, trying again in %s seconds...", 108 | function.__name__, err.response['Error']['Code'], delay 109 | ) 110 | time.sleep(delay) 111 | m_attempts -= 1 112 | last_exception = err 113 | 114 | else: 115 | logger.warning("Error wasn't found in retry raising error:%s", err) 116 | raise err 117 | 118 | logger.error("Was not successfully able to complete the request after %s attempts", max_attempts) 119 | raise last_exception 120 | 121 | return wrapper 122 | 123 | return retry_decorator 124 | 125 | 126 | def validate_ecs_run_task_info(ecs_info: dict): 127 | """Used to validate that SQS Message contains all necessary information to run the ECS Task 128 | 129 | Args: 130 | ecs_info (int): ECS Information 131 | 132 | Returns: 133 | :obj:`json`: Returns json object (dict) of the ECS information 134 | """ 135 | logger.info("Validating SQS Message contains all the needed information to execution ECS Task") 136 | required_net_keys = { 137 | "subnets": ecs_info.get('networkConfiguration', {}).get('awsvpcConfiguration', {}).get('subnets'), 138 | "securityGroups": ecs_info.get('networkConfiguration', {}).get('awsvpcConfiguration', {}).get('securityGroups') 139 | } 140 | required_definition_keys = { 141 | "container_name": ecs_info.get('overrides', {}).get('containerOverrides', [{}])[0].get('name') 142 | } 143 | 144 | ecs = boto3.client('ecs') 145 | logger.info(f"ecs_info:{ecs_info}") 146 | 147 | if None in required_net_keys.values(): 148 | logger.warning("***** REQUIRED NETWORK KEYS NOT PASSED TO RUN TASK *****") 149 | logger.warning(required_net_keys) 150 | 151 | if ecs_info.get('service') and ecs_info.get('cluster'): 152 | response = ecs.describe_services( 153 | cluster=ecs_info['cluster'], 154 | services=[ecs_info['service']] 155 | ) 156 | 157 | logger.debug(response) 158 | if not required_net_keys['subnets'] and not required_net_keys['securityGroups']: 159 | logger.info(f"Updating Networking Configuration:{response['services'][0]['networkConfiguration']}") 160 | ecs_info['networkConfiguration'] = response['services'][0]['networkConfiguration'] 161 | 162 | else: 163 | if not required_net_keys['subnets']: 164 | subnets = response['services'][0]['networkConfiguration']['awsvpcConfiguration']['subnets'] 165 | logger.info(f"Updating Networking Subnet Configuration:{subnets}") 166 | ecs_info['networkConfiguration'] = {"awsvpcConfiguration": {"subnets": subnets}} 167 | 168 | if not required_net_keys['securityGroups']: 169 | security_groups = response['services'][0]['networkConfiguration']['awsvpcConfiguration']['securityGroups'] 170 | logger.info(f"Updating Networking securityGroups Configuration:{security_groups}") 171 | ecs_info['networkConfiguration'] = {"awsvpcConfiguration": {"securityGroups": security_groups}} 172 | 173 | else: 174 | raise Exception( 175 | "ECS Service and/or Cluster Names did not get passed into ECS Information for " 176 | "validate_ecs_run_task_info" 177 | ) 178 | 179 | if None in required_definition_keys.values(): 180 | logger.warning("***** REQUIRED DEFINITION KEYS NOT PASSED TO RUN TASK *****") 181 | logger.warning(required_definition_keys) 182 | 183 | response = ecs.describe_task_definition( 184 | taskDefinition=ecs_info['taskDefinition'] 185 | ) 186 | logger.debug(response) 187 | 188 | if "overrides" not in ecs_info: 189 | ecs_info['overrides'] = {} 190 | 191 | if "containerOverrides" not in ecs_info['overrides']: 192 | ecs_info['overrides']['containerOverrides'] = [{}] 193 | 194 | container_name = response['taskDefinition']['containerDefinitions'][0]['name'] 195 | ecs_info['overrides']['containerOverrides'][0]['name'] = container_name 196 | 197 | # Cleanup of ecs run task dictionary 198 | logger.info("Removing Service from ecs_info dict") 199 | del ecs_info['service'] 200 | logger.info(f"ecs_parameters:{ecs_info}") 201 | 202 | return ecs_info 203 | 204 | 205 | # Max 14 min retry window b/c of Lambdas timeout window 206 | @retry( 207 | max_attempts=44, 208 | delay=randint(15, 20), 209 | error_code=["ThrottlingException", "ServerException", "RequestLimitExceeded"], # InvalidParameterException 210 | error_message="You've reached the limit on the number of tasks you can run concurrently" 211 | ) 212 | def ecs_run_task(**ecs_parameters): 213 | """Runs an ECS Task 214 | 215 | Args: 216 | ecs_parameters (dict): Data pass through from SQS to ECS run task 217 | 218 | Returns: 219 | N/A 220 | """ 221 | logger.info("Attempting to start ECS Task") 222 | ecs = boto3.client('ecs') 223 | 224 | response = ecs.run_task(**ecs_parameters) 225 | logger.info(response) 226 | 227 | # If received a failure within the response raise error 228 | if len(response['failures']) > 0: 229 | raise Exception(response['failures'][0]['reason']) 230 | 231 | logger.info("Successfully started container using command:%s", ecs_parameters) 232 | 233 | 234 | def send_message_to_queue(message_body): 235 | """Sends message to SQS Queue if a time out or limit has occurred to allow for more time for cool down if a 236 | large number of tasks are being executed in parallel. 237 | 238 | Args: 239 | message_body (str): ECS Cluster Name 240 | 241 | Returns: 242 | N/A 243 | """ 244 | logger.info("Pushing message:%s back on the queue.", message_body) 245 | sqs = boto3.client('sqs') 246 | response = sqs.send_message( 247 | QueueUrl=os.environ['QUEUE_URL'], 248 | MessageBody=json.dumps(message_body) 249 | ) 250 | logger.info(response) 251 | -------------------------------------------------------------------------------- /src/main.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | import json 5 | import logging 6 | import os 7 | from random import randrange 8 | from helper import ecs_run_task, send_failure_email, send_message_to_queue, validate_ecs_run_task_info 9 | 10 | logging.basicConfig() 11 | logger = logging.getLogger() 12 | logging.getLogger("botocore").setLevel(logging.ERROR) 13 | logger.setLevel(logging.INFO) 14 | 15 | retry_error_codes = ["ThrottlingException", "ServerException", "RequestLimitExceeded", "InvalidParameterException"] 16 | sqs_retry_limit = int(os.environ['SQS_RETRY_LIMIT']) 17 | 18 | 19 | def lambda_handler(event, context): 20 | print(json.dumps(event)) 21 | sqs_retries = 0 # Priming data 22 | 23 | # If event is from a CloudWatch Rule 24 | if event.get('detail'): 25 | subnets = [] 26 | logger.info('****** Found Event *******') 27 | overrides = event['detail']['overrides']['containerOverrides'][0] 28 | container_info = event['detail']['containers'][0] 29 | attachment_details = event['detail']['attachments'][0]['details'] 30 | 31 | logger.info("Gathering data from event") 32 | for att_detail in attachment_details: 33 | for value in att_detail.values(): 34 | if "subnet-" in value: 35 | subnets.append(value) 36 | 37 | logger.info("Building body dictionary") 38 | body = { 39 | 'ECS': { 40 | "cluster": event['detail']['clusterArn'].split('/')[1], 41 | "subnets": subnets, 42 | "cpu": event['detail']['cpu'], 43 | "memory": event['detail']['memory'], 44 | "command": overrides['command'], 45 | "environment": overrides['environment'], 46 | "container_name": container_info['name'], 47 | "reference_id": f"{container_info['taskArn'].split('/')[1]}-{randrange(10)}", 48 | "task_def": event['detail']['taskDefinitionArn'].split('/')[1].split(':')[0], 49 | "startedBy": "CloudWatch Rules State Change to STOPPED", 50 | "security_groups": [os.environ['ECS_SECURITY_GROUP']] 51 | } 52 | } 53 | 54 | logger.info(f"body:{body}") 55 | 56 | # If event is from SQS 57 | elif event.get('Records'): 58 | record = event['Records'][0] 59 | body = json.loads(record['body'].replace("\'", "\"")) 60 | body['ECS']['startedBy'] = record['eventSourceARN'] 61 | 62 | try: 63 | # Add SQS retry to message body if not already there 64 | if not body.get('SQS_Retries'): 65 | logger.info("Setting up SQS Retry count") 66 | body['SQS_Retries'] = 0 67 | 68 | sqs_retries = body['SQS_Retries'] 69 | 70 | ecs_info = validate_ecs_run_task_info( 71 | ecs_info=body['ECS'] 72 | ) 73 | 74 | ecs_run_task(**ecs_info) 75 | 76 | except KeyError as keyerr: 77 | logger.debug("***** KeyError ******") 78 | send_failure_email(error_message=keyerr, add_message=event) 79 | logger.error(keyerr, exc_info=True) 80 | 81 | except Exception as err: 82 | logger.debug("***** Exception ******") 83 | if hasattr(err, 'response'): 84 | if err.response['Error']['Code'] in retry_error_codes and \ 85 | sqs_retries < sqs_retry_limit: 86 | body['SQS_Retries'] += 1 87 | send_message_to_queue(message_body=body) 88 | return 89 | 90 | elif str(err) == "You've reached the limit on the number of tasks you can run concurrently" and \ 91 | sqs_retries < sqs_retry_limit: 92 | body['SQS_Retries'] += 1 93 | send_message_to_queue(message_body=body) 94 | return 95 | 96 | send_failure_email(error_message=err, add_message=event) 97 | logger.error(err, exc_info=True) 98 | -------------------------------------------------------------------------------- /src/requirements.txt: -------------------------------------------------------------------------------- 1 | requests==2.32.0 2 | cfn-flip==1.2.2 -------------------------------------------------------------------------------- /test/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | PROJECT_PATH = os.getcwd() 4 | SOURCE_PATH = os.path.join( 5 | PROJECT_PATH, "src" 6 | ) 7 | sys.path.append(SOURCE_PATH) -------------------------------------------------------------------------------- /test/local_send_message_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | import logging 5 | import json 6 | import time 7 | import yaml 8 | import boto3 9 | 10 | logging.basicConfig() 11 | logger = logging.getLogger() 12 | logging.getLogger("botocore").setLevel(logging.ERROR) 13 | logger.setLevel(logging.INFO) 14 | 15 | NUMBER_OF_MESSAGES = 1 16 | 17 | 18 | def run_spinup_test(cluster, family): 19 | _max_concurrency = 0 20 | concurrency = 1 # priming the data 21 | 22 | # Watch the tasks run 23 | ecs = boto3.client('ecs') 24 | _response = ['priming_data'] 25 | while concurrency != 0: 26 | concurrency = 0 27 | paginator = ecs.get_paginator("list_tasks") 28 | for page in paginator.paginate( 29 | cluster=cluster, 30 | family=family 31 | ): 32 | concurrency = concurrency + len(page['taskArns']) 33 | 34 | logger.debug(_response) 35 | print("*******************************************") 36 | print(f"* {concurrency} Containers Left") 37 | print("*******************************************") 38 | if concurrency > _max_concurrency: 39 | _max_concurrency = concurrency 40 | 41 | time.sleep(5) 42 | 43 | return _max_concurrency 44 | 45 | 46 | if __name__ == "__main__": 47 | COUNT = 0 48 | sqs = boto3.client('sqs') 49 | sts = boto3.client('sts') 50 | 51 | # Needed for SQS 52 | account_id = sts.get_caller_identity()["Account"] 53 | 54 | configFilePath = "test/local_send_message_test.yaml" 55 | with open(configFilePath, 'r') as f: 56 | config = yaml.safe_load(f) 57 | 58 | # Build message body 59 | message_body = dict() 60 | queue_url = f"https://sqs.{config['region']}.amazonaws.com/{account_id}/{config['queue_name']}" 61 | 62 | # Create messages 63 | while COUNT < NUMBER_OF_MESSAGES: 64 | environment = [ 65 | {"name": "s3_image_id", "value": f"{str(COUNT).zfill(5)}.tiff"}, 66 | {"name": "s3_bucket", "value": f"testapp22-{account_id}-us-east-1"}, 67 | {"name": "sleep_time", "value": "10"} 68 | ] 69 | command = ['sh', '-c', 'sleep 120 && df -h'] 70 | message_body['ECS'] = { 71 | # REQUIRED 72 | "cluster": config['cluster'], 73 | "taskDefinition": config['task_definition'], 74 | "service": config['service'], 75 | "launchType": "FARGATE", 76 | # OPTIONAL 77 | "platformVersion": 'LATEST', 78 | "propagateTags": 'TASK_DEFINITION', 79 | "overrides": { 80 | "containerOverrides": [ 81 | { 82 | "name": config['container_name'], 83 | "command": command, 84 | "environment": environment 85 | } 86 | ], 87 | "cpu": "1024", 88 | "memory": "2048" 89 | }, 90 | # If you'd like to use something other than the default Cluster Subnet and Security Group 91 | # uncomment this section out. 92 | # "networkConfiguration": { 93 | # 'awsvpcConfiguration': { 94 | # 'subnets': config['subnets'], 95 | # 'securityGroups': config['security_group'], 96 | # 'assignPublicIp': "ENABLED" 97 | # } 98 | # }, 99 | "referenceId": "add_reference_id_here", 100 | } 101 | 102 | logger.info('Count:%s', COUNT) 103 | logger.info('Sending message:%s to QueueUrl:%s', str(message_body), queue_url) 104 | response = sqs.send_message( 105 | QueueUrl=queue_url, 106 | MessageBody=json.dumps(message_body) 107 | ) 108 | COUNT += 1 109 | time.sleep(5) 110 | 111 | time.sleep(10) 112 | max_concurrency = run_spinup_test(cluster=config['cluster'], family=config['family']) 113 | print("*******************************************") 114 | print(f"* Total # of Messages: {NUMBER_OF_MESSAGES}") 115 | print(f"* Max Concurrency: {max_concurrency}") 116 | print("*******************************************") 117 | -------------------------------------------------------------------------------- /test/local_send_message_test.yaml: -------------------------------------------------------------------------------- 1 | queue_name: "EcsRunTask-Queue" 2 | cluster: "container-test01" 3 | service: "container-test01" 4 | task_definition: "container-test01:6" 5 | container_name: "container-test01-development" 6 | family: "container-test01" 7 | region: "us-east-1" 8 | 9 | # If you'd like to use something other than the default Cluster Subnet and Security Group 10 | # uncomment this section out and enter in the appropriate Ids. Also uncomment the networkConfiguration 11 | # section within then test/local_send_message_test.py 12 | 13 | #subnets: 14 | # - "subnet-1234567890abcdef0" 15 | # - "subnet-abcdef01234567890" 16 | #security_group: 17 | # - "sg-021345abcdef6789" -------------------------------------------------------------------------------- /test/sqs_messages/test_minimum_requirements.json: -------------------------------------------------------------------------------- 1 | { 2 | "cluster": "container-test01", 3 | "taskDefinition": "container-test01:6", 4 | "service": "container-test01", 5 | "launchType": "FARGATE" 6 | } -------------------------------------------------------------------------------- /test/sqs_messages/test_network_override.json: -------------------------------------------------------------------------------- 1 | { 2 | "cluster": "container-test01", 3 | "taskDefinition": "container-test01:6", 4 | "service": "container-test01", 5 | "launchType": "FARGATE", 6 | "platformVersion": "LATEST", 7 | "propagateTags": "TASK_DEFINITION", 8 | "overrides": { 9 | "containerOverrides": [ 10 | { 11 | "name": "container-test01-development", 12 | "command": [ 13 | "sh", 14 | "-c", 15 | "sleep 120 && df -h" 16 | ], 17 | "environment": [ 18 | { 19 | "name": "s3_image_id", 20 | "value": "00000.tiff" 21 | }, 22 | { 23 | "name": "s3_bucket", 24 | "value": "testapp22-111122223333-us-east-1" 25 | }, 26 | { 27 | "name": "sleep_time", 28 | "value": "1" 29 | } 30 | ], 31 | "cpu": 2048, 32 | "memory": 2048 33 | } 34 | ] 35 | }, 36 | "networkConfiguration": { 37 | "awsvpcConfiguration": { 38 | "subnets": [ 39 | "subnet-1234567890abcdef0", 40 | "subnet-abcdef01234567890" 41 | ], 42 | "securityGroups": [ 43 | "sg-021345abcdef6789" 44 | ], 45 | "assignPublicIp": "ENABLED" 46 | } 47 | }, 48 | "referenceId": "blah", 49 | "startedBy": "arn:aws:sqs:us-east-1:111122223333:EcsRunTask-Queue" 50 | } -------------------------------------------------------------------------------- /test/sqs_messages/test_no_container_name.json: -------------------------------------------------------------------------------- 1 | { 2 | "cluster": "container-test01", 3 | "taskDefinition": "container-test01:6", 4 | "service": "container-test01", 5 | "launchType": "FARGATE", 6 | "platformVersion": "LATEST", 7 | "propagateTags": "TASK_DEFINITION", 8 | "overrides": { 9 | "containerOverrides": [ 10 | { 11 | "command": [ 12 | "sh", 13 | "-c", 14 | "sleep 120 && df -h" 15 | ], 16 | "environment": [ 17 | { 18 | "name": "s3_image_id", 19 | "value": "00000.tiff" 20 | }, 21 | { 22 | "name": "s3_bucket", 23 | "value": "testapp22-111122223333-us-east-1" 24 | }, 25 | { 26 | "name": "sleep_time", 27 | "value": "1" 28 | } 29 | ], 30 | "cpu": 2048, 31 | "memory": 2048 32 | } 33 | ] 34 | }, 35 | "networkConfiguration": { 36 | "awsvpcConfiguration": { 37 | "subnets": [ 38 | "subnet-1234567890abcdef0", 39 | "subnet-abcdef01234567890" 40 | ], 41 | "securityGroups": [ 42 | "sg-021345abcdef6789" 43 | ], 44 | "assignPublicIp": "ENABLED" 45 | } 46 | }, 47 | "referenceId": "blah", 48 | "startedBy": "arn:aws:sqs:us-east-1:111122223333:EcsRunTask-Queue" 49 | } -------------------------------------------------------------------------------- /test/sqs_messages/test_no_network.json: -------------------------------------------------------------------------------- 1 | { 2 | "cluster": "container-test01", 3 | "taskDefinition": "container-test01:6", 4 | "service": "container-test01", 5 | "launchType": "FARGATE", 6 | "platformVersion": "LATEST", 7 | "propagateTags": "TASK_DEFINITION", 8 | "overrides": { 9 | "containerOverrides": [ 10 | { 11 | "name": "container-test01-development", 12 | "command": [ 13 | "sh", 14 | "-c", 15 | "sleep 120 && df -h" 16 | ], 17 | "environment": [ 18 | { 19 | "name": "s3_image_id", 20 | "value": "00000.tiff" 21 | }, 22 | { 23 | "name": "s3_bucket", 24 | "value": "testapp22-111122223333-us-east-1" 25 | }, 26 | { 27 | "name": "sleep_time", 28 | "value": "1" 29 | } 30 | ], 31 | "cpu": 2048, 32 | "memory": 2048 33 | } 34 | ] 35 | }, 36 | "referenceId": "blah", 37 | "startedBy": "arn:aws:sqs:us-east-1:111122223333:EcsRunTask-Queue" 38 | } -------------------------------------------------------------------------------- /test/sqs_messages/test_no_security_group.json: -------------------------------------------------------------------------------- 1 | { 2 | "cluster": "container-test01", 3 | "taskDefinition": "container-test01:6", 4 | "service": "container-test01", 5 | "launchType": "FARGATE", 6 | "platformVersion": "LATEST", 7 | "propagateTags": "TASK_DEFINITION", 8 | "overrides": { 9 | "containerOverrides": [ 10 | { 11 | "name": "container-test01-development", 12 | "command": [ 13 | "sh", 14 | "-c", 15 | "sleep 120 && df -h" 16 | ], 17 | "environment": [ 18 | { 19 | "name": "s3_image_id", 20 | "value": "00000.tiff" 21 | }, 22 | { 23 | "name": "s3_bucket", 24 | "value": "testapp22-111122223333-us-east-1" 25 | }, 26 | { 27 | "name": "sleep_time", 28 | "value": "1" 29 | } 30 | ], 31 | "cpu": 2048, 32 | "memory": 2048 33 | } 34 | ] 35 | }, 36 | "networkConfiguration": { 37 | "awsvpcConfiguration": { 38 | "subnets": [ 39 | "subnet-1234567890abcdef0", 40 | "subnet-abcdef01234567890" 41 | ], 42 | "assignPublicIp": "ENABLED" 43 | } 44 | }, 45 | "referenceId": "blah", 46 | "startedBy": "arn:aws:sqs:us-east-1:111122223333:EcsRunTask-Queue" 47 | } -------------------------------------------------------------------------------- /test/test_ecs_run_task.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | import unittest 5 | import json 6 | from src.helper import validate_ecs_run_task_info 7 | 8 | # ------------------- 9 | # FAILING TEMPLATES 10 | # ------------------- 11 | # Minimum Requirements 12 | f = open("test/sqs_messages/test_minimum_requirements.json", "r") 13 | test_minimum_requirements = json.loads(f.read()) 14 | 15 | # Network Override 16 | f = open("test/sqs_messages/test_network_override.json", "r") 17 | test_network_override = json.loads(f.read()) 18 | 19 | # No Container Name 20 | f = open("test/sqs_messages/test_no_container_name.json", "r") 21 | test_no_container_name = json.loads(f.read()) 22 | 23 | # No Network 24 | f = open("test/sqs_messages/test_no_network.json", "r") 25 | test_no_network = json.loads(f.read()) 26 | 27 | # No Security Group 28 | f = open("test/sqs_messages/test_no_security_group.json", "r") 29 | test_no_security_group = json.loads(f.read()) 30 | 31 | f.close() 32 | 33 | 34 | class TestStringMethods(unittest.TestCase): 35 | 36 | # Pass 37 | def test_minimum_requirements(self): 38 | self.assertTrue(validate_ecs_run_task_info(ecs_info=test_minimum_requirements)) 39 | 40 | def test_network_override(self): 41 | self.assertTrue(validate_ecs_run_task_info(ecs_info=test_network_override)) 42 | 43 | def test_no_container_name(self): 44 | self.assertTrue(validate_ecs_run_task_info(ecs_info=test_no_container_name)) 45 | 46 | def test_no_network(self): 47 | self.assertTrue(validate_ecs_run_task_info(ecs_info=test_no_network)) 48 | 49 | def test_no_security_group(self): 50 | self.assertTrue(validate_ecs_run_task_info(ecs_info=test_no_security_group)) 51 | 52 | 53 | if __name__ == '__main__': 54 | unittest.main() 55 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | skipsdist = true 3 | 4 | [testenv] 5 | deps = 6 | -rrequirements_validate_scan.txt 7 | 8 | setenv = 9 | PYTHONPATH = {toxinidir} 10 | 11 | commands = 12 | coverage run -m pytest . --cov=. --cov-report xml --cov-report html --junitxml=junitxml.xml 13 | --------------------------------------------------------------------------------