├── __init__.py ├── tests ├── accounts2.csv ├── accounts.csv └── test_scripts.py ├── NOTICE ├── CODE_OF_CONDUCT.md ├── README.md ├── EnableDetective.yaml ├── CONTRIBUTING.md ├── LICENSE ├── disableDetective.py └── enableDetective.py /__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/accounts2.csv: -------------------------------------------------------------------------------- 1 | 2 | 3 | -------------------------------------------------------------------------------- /NOTICE: -------------------------------------------------------------------------------- 1 | Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | -------------------------------------------------------------------------------- /tests/accounts.csv: -------------------------------------------------------------------------------- 1 | 123456789012,random@gmail.com 2 | 000012345678,email@gmail.com 3 | 12345,wrongaccount@gmail.com 4 | 555555555555 ,test5@gmail.com 5 | 111111111111, test1@gmail.com 6 | 222222222222, test2@gmail.com 7 | 333333333333, test3@gmail.com 8 | 9 | 1234987b 10 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # amazon-detective-multiaccount-scripts 2 | 3 | Amazon Detective provides a set of open-source Python scripts in this repository. The scripts require Python 3. 4 | 5 | You can use these to perform the following tasks: 6 | * Enable Detective for an administrator account across Regions. When you enable Detective, you can assign tag values to the behavior graph. 7 | * Add member accounts to an administrator account's behavior graphs across Regions. 8 | * Optionally send invitation emails to the member accounts. You can also configure the request to not send invitation emails. 9 | * Remove member accounts from an administrator account's behavior graphs across Regions. 10 | * Disable Detective for an administrator account across Regions. When an administrator account disables Detective, the administrator account's behavior graph in each Region is disabled. 11 | 12 | For more information on how to use these scripts, see [Using the Amazon Detective Python scripts](https://docs.aws.amazon.com/detective/latest/adminguide/detective-github-scripts.html) 13 | 14 | ## Contributing to this project 15 | 16 | ### Running tests 17 | 18 | ``` 19 | # Install requirements 20 | pip3 install boto3 pytest 21 | 22 | # In the tests/ directory... 23 | pytest -s 24 | ``` 25 | -------------------------------------------------------------------------------- /EnableDetective.yaml: -------------------------------------------------------------------------------- 1 | AWSTemplateFormatVersion: 2010-09-09 2 | Description: Creates a new role to allow an administrator account to enable and manage Detective. 3 | 4 | Parameters: 5 | AdministratorAccountId: 6 | Type: String 7 | Description: AWS Account Id of the administrator account (the account in which will recieve Detective findings from member accounts). 8 | MaxLength: 12 9 | MinLength: 12 10 | CreateInstanceRole: 11 | Type: String 12 | Description: Select Yes to create an EC2 instance role that can be attached to an instnace in the Master account which will allow the instance to assume the exection role. Select No if you plan to run the script locally or are creating the stack in a member account. 13 | AllowedValues: ["Yes", "No"] 14 | Conditions: 15 | CreateInstanceRole: !Equals [!Ref CreateInstanceRole, "Yes"] 16 | Resources: 17 | ExecutionRole: 18 | Type: AWS::IAM::Role 19 | Properties: 20 | RoleName: ManageDetective 21 | AssumeRolePolicyDocument: 22 | Version: 2012-10-17 23 | Statement: 24 | - Effect: Allow 25 | Principal: 26 | AWS: 27 | - !Ref AdministratorAccountId 28 | Action: 29 | - sts:AssumeRole 30 | Path: / 31 | ManagedPolicyArns: 32 | - arn:aws:iam::aws:policy/AmazonDetectiveFullAccess 33 | InstanceRole: 34 | Type: AWS::IAM::Role 35 | Condition: CreateInstanceRole 36 | Properties: 37 | RoleName: ManageDetectiveInstanceRole 38 | AssumeRolePolicyDocument: 39 | Version: "2012-10-17" 40 | Statement: 41 | - 42 | Effect: "Allow" 43 | Principal: 44 | Service: 45 | - "ec2.amazonaws.com" 46 | Action: 47 | - "sts:AssumeRole" 48 | Policies: 49 | - 50 | PolicyName: ManageDetectivePolicy 51 | PolicyDocument: 52 | Version: "2012-10-17" 53 | Statement: 54 | - 55 | Effect: "Allow" 56 | Action: "sts:AssumeRole" 57 | Resource: !Join ["", ["arn:aws:iam::*:role/",!Ref ExecutionRole]] 58 | InstanceProfile: 59 | Type: AWS::IAM::InstanceProfile 60 | Condition: CreateInstanceRole 61 | Properties: 62 | Path: / 63 | Roles: 64 | - !Ref InstanceRole 65 | InstanceProfileName: EnableDetective 66 | 67 | -------------------------------------------------------------------------------- /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 *master* 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 | 61 | We may ask you to sign a [Contributor License Agreement (CLA)](http://en.wikipedia.org/wiki/Contributor_License_Agreement) for larger changes. 62 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /disableDetective.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | """ python3 disableDetective.py --master_account 555555555555 --assume_role detectiveAdmin --disabled_regions us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1 --input_file accounts.csv 4 | """ 5 | __author__ = "Ryan Nolette" 6 | __copyright__ = "Amazon 2020" 7 | __credits__ = ["Ryan Nolette", 8 | "https://github.com/sonofagl1tch"] 9 | __license__ = "Apache" 10 | __version__ = "1.0.1" 11 | __maintainer__ = "Ryan Nolette" 12 | __email__ = "detective-demo-requests@amazon.com" 13 | __status__ = "Production" 14 | 15 | import argparse 16 | import itertools 17 | import logging 18 | import re 19 | import sys 20 | import typing 21 | import boto3 22 | import botocore.exceptions 23 | 24 | FORMAT = '%(asctime)s - %(levelname)s - %(message)s' 25 | logging.basicConfig(level=logging.INFO, stream=sys.stdout, format=FORMAT) 26 | 27 | 28 | def setup_command_line(args = None) -> argparse.Namespace: 29 | """ 30 | Configures and reads command line arguments. 31 | 32 | Returns: 33 | An argparse.Namespace object containing parsed arguments. 34 | 35 | Raises: 36 | argpare.ArgumentTypeError if an invalid value is used for 37 | master_account argument. 38 | """ 39 | def _master_account_type(val: str, pattern: str = r'[0-9]{12}'): 40 | if not re.match(pattern, val): 41 | raise argparse.ArgumentTypeError 42 | return val 43 | 44 | # Setup command line arguments 45 | parser = argparse.ArgumentParser(description=('Unlink AWS Accounts from a central ' 46 | 'Detective Account, or delete the entire Detective graph.')) 47 | parser.add_argument('--master_account', type=_master_account_type, 48 | required=True, 49 | help="AccountId for Central AWS Account.") 50 | parser.add_argument('--input_file', type=argparse.FileType('r'), 51 | help=('Path to CSV file containing the list of ' 52 | 'account IDs and Email addresses. ' 53 | 'This does not need to be provided if you use the delete_graph flag.')) 54 | parser.add_argument('--assume_role', type=str, required=True, 55 | help="Role Name to assume in each account.") 56 | parser.add_argument('--delete_graph', action='store_true', 57 | help=('Delete the master Detective graph. ' 58 | 'If not provided, you must provide an input file.')) 59 | parser.add_argument('--disabled_regions', type=str, 60 | help=('Regions to disable Detective. If not specified, ' 61 | 'all available regions disabled.')) 62 | args = parser.parse_args(args) 63 | if not args.delete_graph and not args.input_file: 64 | raise parser.error("Either an input file or the delete_graph flag should be provided.") 65 | 66 | return args 67 | 68 | 69 | def read_accounts_csv(input_file: typing.IO) -> typing.Dict: 70 | """ 71 | Parses contents from the CSV file containing the accounts and email addreses. 72 | 73 | Args: 74 | input_file: A file object to read CSV data from. 75 | 76 | Returns: 77 | A dictionary where the key is account ID and value is email address. 78 | """ 79 | account_re = re.compile(r'[0-9]{12}') 80 | aws_account_dict = {} 81 | 82 | if not input_file: 83 | return aws_account_dict 84 | 85 | for acct in input_file.readlines(): 86 | split_line = acct.strip().split(',') 87 | 88 | if len(split_line) != 2: 89 | logging.exception(f'Unable to process line: {acct}.') 90 | continue 91 | 92 | account_number, email = split_line 93 | if not account_re.match(account_number): 94 | logging.error( 95 | f'Invalid account number {account_number}, skipping.') 96 | continue 97 | 98 | aws_account_dict[account_number.strip()] = email.strip() 99 | 100 | return aws_account_dict 101 | 102 | 103 | def get_regions(session: boto3.Session, user_regions: str) -> typing.List[str]: 104 | """ 105 | Get AWS regions to disable Detective from. 106 | 107 | Args: 108 | session: boto3 session. 109 | args_region: User specificied regions. 110 | 111 | Returns: 112 | A list of the region names to disable Detective from. 113 | """ 114 | detective_regions = [] 115 | if user_regions: 116 | detective_regions = user_regions.split(',') 117 | logging.info( 118 | f'disabling members in these regions: {detective_regions}') 119 | else: 120 | detective_regions = session.get_available_regions('detective') 121 | logging.info( 122 | f'disabling members in all available Detective regions {detective_regions}') 123 | return detective_regions 124 | 125 | 126 | def assume_role(aws_account_number: str, role_name: str) -> boto3.Session: 127 | """ 128 | Assumes the provided role in each account and returns a Detective client. 129 | 130 | Args: 131 | - aws_account_number: AWS Account Number 132 | - role_name: Role to assume in target account 133 | 134 | Returns: 135 | Detective client in the specified AWS Account and Region 136 | """ 137 | try: 138 | # Beginning the assume role process for account 139 | sts_client = boto3.client('sts') 140 | 141 | # Get the current partition 142 | partition = sts_client.get_caller_identity()['Arn'].split(":")[1] 143 | 144 | response = sts_client.assume_role( 145 | RoleArn='arn:{}:iam::{}:role/{}'.format( 146 | partition, 147 | aws_account_number, 148 | role_name 149 | ), 150 | RoleSessionName='EnableDetective' 151 | ) 152 | # Storing STS credentials 153 | session = boto3.Session( 154 | aws_access_key_id=response['Credentials']['AccessKeyId'], 155 | aws_secret_access_key=response['Credentials']['SecretAccessKey'], 156 | aws_session_token=response['Credentials']['SessionToken'] 157 | ) 158 | except Exception as e: 159 | logging.exception(f'exception: {e}') 160 | 161 | logging.info(f"Assumed session for {aws_account_number}.") 162 | 163 | return session 164 | 165 | 166 | def get_graphs(d_client: botocore.client.BaseClient) -> typing.List[str]: 167 | """ 168 | Get graphs in a specified region. 169 | 170 | Args: 171 | - d_client: Detective boto3 client generated from the master session. 172 | 173 | Returns: 174 | List of graph Arns. 175 | """ 176 | try: 177 | response = d_client.list_graphs() 178 | except botocore.exceptions.EndpointConnectionError: 179 | logging.exception(f'exception: {e}') 180 | return [] 181 | 182 | # use .get function to avoid KeyErrors when a dictionary key doesn't exist 183 | # it returns an empty list instead. 184 | # map iterates over all elements under a list and applied a function to them, 185 | # in this specific case, the element 'Arn' is extracted from the dictionary 186 | # (graphlist is a list of dictionaries) 187 | return [x['Arn'] for x in response.get('GraphList', [])] 188 | 189 | 190 | def get_members(d_client: botocore.client.BaseClient, graphs: typing.List[str]) ->\ 191 | (typing.Dict[str, typing.Set[str]], typing.Dict[str, typing.Set[str]]): 192 | """ 193 | Get member accounts for all behaviour graphs in a region. 194 | 195 | Args: 196 | - d_client: Detective boto3 client generated from the master session. 197 | - graphs: List of graphs arns 198 | 199 | Returns: 200 | Two dictionaries: one with all account ids, other with the ones pending to accept 201 | the invitation. 202 | """ 203 | try: 204 | # itertools.tee creates two independent iterators from a single one. This way 205 | # we can iterate the iterator twice: one to return all elements and other to return 206 | # the ones pending to be invited. 207 | #### 208 | # check the value of NextToken in the response. if it is non-null, pass it back into a subsequent list_members call (and keep doing this until a null token is returned) 209 | def _master_memberList(g: str) -> typing.List[typing.Dict]: 210 | # create a list to append the member accounts 211 | memberAccounts = [] 212 | # create a dictionary for the nextToken from each call 213 | tokenTracker = {} 214 | # loop through list_members call results and take action for each returned result 215 | while True: 216 | # list_members of graph "g" and return the first 100 results 217 | members = d_client.list_members( 218 | GraphArn=g, MaxResults=100, **tokenTracker) 219 | # add the returned list members to the list 220 | memberAccounts.extend(members['MemberDetails']) 221 | # if the returned results have a "NextToken" key then use it to query again 222 | if 'NextToken' in members: 223 | tokenTracker['NextToken'] = members['NextToken'] 224 | # if the returned results do not have a "NextToken" key then exit the loop 225 | else: 226 | break 227 | # return members list. 228 | # The return statement doesn't need () 229 | return memberAccounts 230 | except Exception as e: 231 | logging.exception(f'exception when getting memebers: {e}') 232 | # iterate through each list and return results 233 | all_ac, pending = itertools.tee((g, _master_memberList(g)) 234 | for g in graphs) 235 | return ({g: {x['AccountId'] for x in v} for g, v in all_ac}, 236 | {g: {x['AccountId'] for x in v if x['Status'] == 'INVITED'} for g, v in pending}) 237 | 238 | 239 | def delete_members(d_client: botocore.client.BaseClient, graph_arn: str, 240 | account_dict: typing.Dict[str, str]) -> typing.Set[str]: 241 | """ 242 | delete member accounts for all accounts in the csv that are not present in the graph member set. 243 | 244 | Args: 245 | - d_client: Detective boto3 client generated from the master session. 246 | - graph_arn: Graph to add members to. 247 | - account_dict: Accounts read from the CSV input file. 248 | 249 | Returns: 250 | Set with the IDs of the successfully deleted accounts. 251 | """ 252 | try: 253 | accountIDs = list(account_dict.keys()) 254 | response = d_client.delete_members(GraphArn=graph_arn, 255 | AccountIds=accountIDs) 256 | for error in response['UnprocessedAccounts']: 257 | logging.exception(f'Could not delete member for account {error["AccountId"]} in ' 258 | f'graph {graph_arn}: {error["Reason"]}') 259 | except e: 260 | logging.error(f'error when deleting member: {e}') 261 | 262 | def chunked(it, size): 263 | it = iter(it) 264 | while True: 265 | p = tuple(itertools.islice(it, size)) 266 | if not p: 267 | break 268 | yield p 269 | 270 | if __name__ == '__main__': 271 | args = setup_command_line() 272 | aws_account_dict = read_accounts_csv(args.input_file) 273 | 274 | # making sure that we either have an account list to delete from graphs or the delete_graph flag is provided. 275 | if len(list(aws_account_dict.keys())) == 0 and not args.delete_graph: 276 | logging.error("The delete_graph flag was False while the provided account list is empty. "\ 277 | "Please check your inputs and re-run the script.") 278 | exit(1) 279 | 280 | try: 281 | session = boto3.session.Session() 282 | detective_regions = get_regions(session, args.disabled_regions) 283 | master_session = assume_role(args.master_account, args.assume_role) 284 | except NameError as e: 285 | # logging.exception prints the full traceback, logging.error just prints the error message. 286 | # In this case, the NameError is handled in a specific way: it is an expected exception 287 | # and the traceback doesn't add any value to the error message. 288 | logging.error(f'Master account is not defined: {e.args}') 289 | except Exception as e: 290 | # in this case, there has been an unhandled exception, something that wasn't estimated. 291 | # Having the error traceback helps us know what happened and help us find a solution 292 | # for the bug. The code should never arrive to this except clause, but if it does 293 | # we want as much information as possible and that's why we use logging.traceback. 294 | # In this case the traceback adds LOTS of value. 295 | logging.exception(f'error creating session {e.args}') 296 | 297 | #Chunk the list of accounts in the .csv into batches of 50 due to the API limitation of 50 accounts per invokation 298 | for chunk in chunked(aws_account_dict.items(), 50): 299 | 300 | for region in detective_regions: 301 | try: 302 | d_client = master_session.client('detective', region_name=region) 303 | graphs = get_graphs(d_client) 304 | if not graphs: 305 | logging.info(f'Amazon Detective has already been disabled in {region}') 306 | else: 307 | logging.info(f'Disabling Amazon Detective in region {region}') 308 | 309 | try: 310 | for graph in graphs: 311 | if not args.delete_graph: 312 | delete_members(d_client, graph, chunk) 313 | else: 314 | d_client.delete_graph(graph) 315 | except NameError as e: 316 | logging.error(f'account is not defined: {e}') 317 | except Exception as e: 318 | logging.exception(f'{e}') 319 | 320 | except NameError as e: 321 | logging.error(f'account is not defined: {e}') 322 | except Exception as e: 323 | logging.exception(f'error with region {region}: {e}') 324 | -------------------------------------------------------------------------------- /tests/test_scripts.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import sys 3 | sys.path.append("..") 4 | 5 | import enableDetective 6 | import disableDetective 7 | import pytest 8 | from unittest.mock import Mock 9 | 10 | ### 11 | # The purpose of this test is to make sure our internal method _master_account_type() deals with master accounts correctly and the rest of the inputs are parsed correctly as well in enableDetective.py. 12 | ### 13 | def test_setup_command_line_enableDetective(): 14 | 15 | args = enableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--enabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 16 | assert args.master_account == '555555555555' 17 | assert args.assume_role == 'detectiveAdmin' 18 | assert args.enabled_regions == 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1' 19 | assert args.disable_email == False 20 | 21 | args = enableDetective.setup_command_line(['--master_account', '012345678901', '--assume_role', 'detectiveAdmin', '--input_file', 'accounts.csv']) 22 | assert args.master_account == '012345678901' 23 | assert args.enabled_regions == None 24 | 25 | args = enableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--enabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 26 | assert args.master_account == '000000000001' 27 | assert args.tags == None 28 | 29 | args = enableDetective.setup_command_line("--master_account 123456789012 --assume_role detectiveAdmin --input_file accounts.csv --tags TagKey1=TagValue1,TagKey2=TagValue2,TagKey3=TagValue3,TagKey4=,TagKey5=TagValue5,TagKey6".split(" ")) 30 | assert args.tags == { 31 | "TagKey1": "TagValue1", 32 | "TagKey2": "TagValue2", 33 | "TagKey3": "TagValue3", 34 | "TagKey4": "", 35 | "TagKey5": "TagValue5", 36 | "TagKey6": "", 37 | } 38 | 39 | args = enableDetective.setup_command_line(['--disable_email', '--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--input_file', 'accounts.csv']) 40 | assert args.disable_email == True 41 | 42 | # Wrong master account 43 | # The internal function _master_account_type() should raise argparse.ArgumentTypeError, however this exception gets supressed by argparse, and SystemExit is raised instead. 44 | with pytest.raises(SystemExit): 45 | enableDetective.setup_command_line(['--master_account', '12345', '--assume_role', 'detectiveAdmin', '--enabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 46 | 47 | # Non existent input file 48 | with pytest.raises(SystemExit): 49 | enableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--enabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts1.csv']) 50 | 51 | 52 | ### 53 | # The purpose of this test is to make sure our internal method _master_account_type() deals with master accounts correctly and the rest of the inputs are parsed correctly as well in disableDetective.py. 54 | ### 55 | def test_setup_command_line_disableDetective(): 56 | 57 | args = disableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 58 | assert args.master_account == '555555555555' 59 | assert args.assume_role == 'detectiveAdmin' 60 | assert args.disabled_regions == 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1' 61 | 62 | args = disableDetective.setup_command_line(['--master_account', '012345678901', '--assume_role', 'detectiveAdmin', '--input_file', 'accounts.csv']) 63 | assert args.master_account == '012345678901' 64 | assert args.disabled_regions == None 65 | 66 | args = disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 67 | assert args.master_account == '000000000001' 68 | 69 | # Wrong master account 70 | # The internal function _master_account_type() should raise argparse.ArgumentTypeError, however this exception gets supressed by argparse, and SystemExit is raised instead. 71 | with pytest.raises(SystemExit): 72 | disableDetective.setup_command_line(['--master_account', '12345', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 73 | 74 | # Non existent input file and no delete_graph flag 75 | with pytest.raises(SystemExit): 76 | disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts1.csv']) 77 | 78 | # Non existent input file with delete_graph flag provided 79 | with pytest.raises(SystemExit): 80 | disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts1.csv', '--delete_graph']) 81 | 82 | # No input file provided and no delete_graph flag 83 | with pytest.raises(SystemExit): 84 | disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1']) 85 | 86 | # No input file provided but delete_graph is set 87 | args = disableDetective.setup_command_line(['--master_account', '000000000001', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--delete_graph']) 88 | assert args.input_file == None 89 | assert args.delete_graph == True 90 | 91 | ### 92 | # The purpose of this test is to make sure we read accounts and emails correctly from the input .csv file in enableDetective.py 93 | ### 94 | def test_read_accounts_csv_enableDetective(): 95 | args = enableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--enabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 96 | 97 | accounts_dict = enableDetective.read_accounts_csv(args.input_file) 98 | 99 | assert len(accounts_dict.keys()) == 6 100 | assert accounts_dict == {"123456789012":"random@gmail.com", "000012345678":"email@gmail.com", "555555555555":"test5@gmail.com", "111111111111":"test1@gmail.com", "222222222222":"test2@gmail.com", "333333333333":"test3@gmail.com"} 101 | 102 | ### 103 | # The purpose of this test is to make sure we read accounts and emails correctly from the input .csv file in disableDetective.py 104 | ### 105 | def test_read_accounts_csv_disableDetective(): 106 | # a test case where an input file is provided, although some of the lines are not correct 107 | args = disableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts.csv']) 108 | 109 | accounts_dict = disableDetective.read_accounts_csv(args.input_file) 110 | 111 | assert len(accounts_dict.keys()) == 6 112 | assert accounts_dict == {"123456789012":"random@gmail.com", "000012345678":"email@gmail.com", "555555555555":"test5@gmail.com", "111111111111":"test1@gmail.com", "222222222222":"test2@gmail.com", "333333333333":"test3@gmail.com"} 113 | 114 | # a test case where no input file is provided 115 | args = disableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--delete_graph']) 116 | assert args.input_file == None 117 | accounts_dict = disableDetective.read_accounts_csv(args.input_file) 118 | assert accounts_dict == {} 119 | 120 | # a test case where an empty input file is provided, along with delete_graph. This is not an error, although clients should not run with this kind of input 121 | args = disableDetective.setup_command_line(['--master_account', '555555555555', '--assume_role', 'detectiveAdmin', '--disabled_regions', 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1', '--input_file', 'accounts2.csv', '--delete_graph']) 122 | accounts_dict = disableDetective.read_accounts_csv(args.input_file) 123 | assert accounts_dict == {} 124 | assert args.delete_graph == True 125 | 126 | ### 127 | # The purpose of this test is to make sure we extract regions correctly in enableDetective.py 128 | ### 129 | def test_get_regions_enableDetective(): 130 | # Regions are provide by user 131 | regions = enableDetective.get_regions(None, 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1') 132 | assert regions == ['us-east-1', 'us-east-2', 'us-west-2', 'ap-northeast-1', 'eu-west-1'] 133 | 134 | # Need to use Boto to get the regions 135 | boto_session = Mock() 136 | boto_session.get_available_regions.return_value = ['us-east-1', 'us-east-2'] 137 | 138 | regions = enableDetective.get_regions(boto_session, None) 139 | assert regions == ['us-east-1', 'us-east-2'] 140 | 141 | 142 | ### 143 | # The purpose of this test is to make sure we extract regions correctly in disableDetective.py 144 | ### 145 | def test_get_regions_disableDetective(): 146 | # Regions are provide by user 147 | regions = disableDetective.get_regions(None, 'us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1') 148 | assert regions == ['us-east-1', 'us-east-2', 'us-west-2', 'ap-northeast-1', 'eu-west-1'] 149 | 150 | # Need to use Boto to get the regions 151 | boto_session = Mock() 152 | boto_session.get_available_regions.return_value = ['us-east-1', 'us-east-2'] 153 | 154 | regions = disableDetective.get_regions(boto_session, None) 155 | assert regions == ['us-east-1', 'us-east-2'] 156 | 157 | ### 158 | # The purpose of this test is to make sure get_members() runs correctly in both enableDetective.py and disableDetective.py 159 | ### 160 | def test_get_members_enableDetective_and_disableDetective(): 161 | d_client = Mock() 162 | 163 | # Test case where we have both active and pending accounts 164 | d_client.list_members.return_value = {"MemberDetails":[{"AccountId":"123456789012", "Status":"ENABLED"}, {"AccountId":"111111111111", "Status":"ENABLED"}, {"AccountId":"222222222222", "Status":"INVITED"}]} 165 | 166 | # Check correctness in enableDetective (with 2 specified graphs) 167 | all_ac, pending = enableDetective.get_members(d_client, ["graph1", "graph2"]) 168 | 169 | assert all_ac == {"graph1": {"123456789012", "111111111111", "222222222222"}, "graph2": {"123456789012", "111111111111", "222222222222"}} 170 | assert pending == {"graph1": {"222222222222"}, "graph2": {"222222222222"}} 171 | 172 | # Check correctness in disableDetective (with 2 specified graphs) 173 | all_ac, pending = disableDetective.get_members(d_client, ["graph1", "graph2"]) 174 | 175 | assert all_ac == {"graph1": {"123456789012", "111111111111", "222222222222"}, "graph2": {"123456789012", "111111111111", "222222222222"}} 176 | assert pending == {"graph1": {"222222222222"}, "graph2": {"222222222222"}} 177 | 178 | # Test case with no pending accounts 179 | d_client.list_members.return_value = {"MemberDetails":[{"AccountId":"111111111111", "Status":"ENABLED"}]} 180 | 181 | # Check correctness in enableDetective (with 1 specified graph) 182 | all_ac, pending = enableDetective.get_members(d_client, ["graph1"]) 183 | 184 | assert all_ac == {"graph1": {"111111111111"}} 185 | assert pending == {"graph1": set()} 186 | 187 | # Check correctness in disableDetective (with 1 specified graph) 188 | all_ac, pending = disableDetective.get_members(d_client, ["graph1"]) 189 | 190 | assert all_ac == {"graph1": {"111111111111"}} 191 | assert pending == {"graph1": set()} 192 | 193 | 194 | def test_create_members(): 195 | d_client = Mock() 196 | 197 | # The existing members and new required members don't overlap. Creation causes no errors. 198 | d_client.create_members.return_value = {'UnprocessedAccounts': {}, 'Members': [{"AccountId": "111111111111"}, {"AccountId": "222222222222"}]} 199 | 200 | created_members = enableDetective.create_members(d_client, "graph1", False, {"333333333333"}, {"111111111111": "1@gmail.com", "222222222222": "2@gmail.com"}) 201 | assert created_members == {"111111111111", "222222222222"} 202 | 203 | # The existing members and new required members overlap by one account. Creation causes no errors. 204 | d_client.create_members.return_value = {'UnprocessedAccounts': {}, 'Members': [{"AccountId": "111111111111"}]} 205 | 206 | created_members = enableDetective.create_members(d_client, "graph1", False, {"222222222222"}, {"111111111111": "1@gmail.com", "222222222222": "2@gmail.com"}) 207 | assert created_members == {"111111111111"} 208 | 209 | # All required members already exist. Creation causes no errors. 210 | d_client.create_members.return_value = {'UnprocessedAccounts': {}, 'Members': []} 211 | 212 | created_members = enableDetective.create_members(d_client, "graph1", False, {"111111111111", "222222222222", "333333333333"}, {"111111111111": "1@gmail.com", "222222222222": "2@gmail.com"}) 213 | assert created_members == set() 214 | 215 | 216 | # The existing members and new required members don't overlap. Creation causes 1 error. 217 | d_client.create_members.return_value = {'UnprocessedAccounts': {"AccountId": "111111111111", "Reason": "some_reason"}, 'Members': [{"AccountId": "222222222222"}]} 218 | 219 | created_members = enableDetective.create_members(d_client, "graph1", False, {"333333333333"}, {"111111111111": "1@gmail.com", "222222222222": "2@gmail.com"}) 220 | assert created_members == {"222222222222"} 221 | 222 | # Test the disabled email flag is properly passed through 223 | enableDetective.create_members(d_client, "with_email", False, {"333333333333"}, {"111111111111": "1@gmail.com"}) 224 | d_client.create_members.assert_called_with(GraphArn="with_email", 225 | Message='Automatically generated invitation', 226 | Accounts=[{'AccountId': '111111111111', 'EmailAddress': '1@gmail.com'}], 227 | DisableEmailNotification=False) 228 | 229 | enableDetective.create_members(d_client, "without_email", False, {"333333333333"}, {"111111111111": "1@gmail.com"}) 230 | d_client.create_members.assert_called_with(GraphArn="without_email", 231 | Message='Automatically generated invitation', 232 | Accounts=[{'AccountId': '111111111111', 'EmailAddress': '1@gmail.com'}], 233 | DisableEmailNotification=False) 234 | 235 | def test_enable_detective(): 236 | d_client = Mock() 237 | 238 | d_client.list_graphs.return_value = {'GraphList': []} 239 | d_client.create_graph.return_value = {'GraphArn': 'fooGraph123'} 240 | 241 | graphs = enableDetective.enable_detective(d_client, "us-east-2") 242 | 243 | assert graphs == ['fooGraph123'] 244 | 245 | 246 | 247 | 248 | -------------------------------------------------------------------------------- /enableDetective.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | """ python3 enableDetective.py --master_account 555555555555 --assume_role detectiveAdmin --enabled_regions us-east-1,us-east-2,us-west-2,ap-northeast-1,eu-west-1 --input_file accounts.csv 4 | """ 5 | __author__ = "Ryan Nolette" 6 | __copyright__ = "Amazon 2020" 7 | __credits__ = ["Ryan Nolette", 8 | "https://github.com/sonofagl1tch"] 9 | __license__ = "Apache" 10 | __version__ = "1.0.1" 11 | __maintainer__ = "Ryan Nolette" 12 | __email__ = "detective-demo-requests@amazon.com" 13 | __status__ = "Production" 14 | 15 | import argparse 16 | import itertools 17 | import logging 18 | import re 19 | import sys 20 | import time 21 | import typing 22 | import boto3 23 | import botocore.exceptions 24 | 25 | FORMAT = '%(asctime)s - %(levelname)s - %(message)s' 26 | logging.basicConfig(level=logging.INFO, stream=sys.stdout, format=FORMAT) 27 | 28 | 29 | def setup_command_line(args = None) -> argparse.Namespace: 30 | """ 31 | Configures and reads command line arguments. 32 | 33 | Returns: 34 | An argparse.Namespace object containing parsed arguments. 35 | 36 | Raises: 37 | argpare.ArgumentTypeError if an invalid value is used for 38 | master_account argument. 39 | """ 40 | def _master_account_type(val: str, pattern: str = r'[0-9]{12}'): 41 | if not re.match(pattern, val): 42 | raise argparse.ArgumentTypeError 43 | return val 44 | 45 | class ParseCommaSeparatedKeyValuePairsAction(argparse.Action): 46 | def __call__(self, parser, namespace, values, option_string=None): 47 | setattr(namespace, self.dest, dict()) 48 | for kv_pairs in values.split(","): 49 | key, _, value = kv_pairs.partition('=') 50 | getattr(namespace, self.dest)[key] = value 51 | 52 | # Setup command line arguments 53 | parser = argparse.ArgumentParser(description=('Link AWS Accounts to central ' 54 | 'Detective Account.')) 55 | parser.add_argument('--master_account', type=_master_account_type, 56 | required=True, 57 | help="AccountId for Central AWS Account.") 58 | parser.add_argument('--input_file', type=argparse.FileType('r'), 59 | required=True, 60 | help=('Path to CSV file containing the list of ' 61 | 'account IDs and Email addresses.')) 62 | parser.add_argument('--assume_role', type=str, required=True, 63 | help="Role Name to assume in each account.") 64 | parser.add_argument('--enabled_regions', type=str, 65 | help=('Regions to enable Detective. If not specified, ' 66 | 'all available regions enabled.')) 67 | parser.add_argument('--disable_email', action='store_true', 68 | help=('Don\'t send emails to the member accounts. Member ' 69 | 'accounts must still accept the invitation before ' 70 | 'they are added to the behavior graph.')) 71 | parser.add_argument('--tags', 72 | action=ParseCommaSeparatedKeyValuePairsAction, 73 | help='Comma-separated list of tag key-value pairs to be added ' 74 | 'to any newly enabled Detective graphs. Values are optional ' 75 | 'and are separated from keys by the equal sign (i.e. \'=\')') 76 | return parser.parse_args(args) 77 | 78 | 79 | def read_accounts_csv(input_file: typing.IO) -> typing.Dict: 80 | """ 81 | Parses contents from the CSV file containing the accounts and email addreses. 82 | 83 | Args: 84 | input_file: A file object to read CSV data from. 85 | 86 | Returns: 87 | A dictionary where the key is account ID and value is email address. 88 | """ 89 | account_re = re.compile(r'[0-9]{12}') 90 | aws_account_dict = {} 91 | 92 | for acct in input_file.readlines(): 93 | split_line = acct.strip().split(',') 94 | 95 | if len(split_line) != 2: 96 | logging.exception(f'Unable to process line: {acct}.') 97 | continue 98 | 99 | account_number, email = split_line 100 | if not account_re.match(account_number): 101 | logging.error( 102 | f'Invalid account number {account_number}, skipping.') 103 | continue 104 | 105 | aws_account_dict[account_number.strip()] = email.strip() 106 | 107 | return aws_account_dict 108 | 109 | 110 | def get_regions(session: boto3.Session, user_regions: str) -> typing.List[str]: 111 | """ 112 | Get AWS regions to enable Detective from. 113 | 114 | Args: 115 | session: boto3 session. 116 | args_region: User specificied regions. 117 | 118 | Returns: 119 | A list of the region names to enable Detective from. 120 | """ 121 | detective_regions = [] 122 | if user_regions: 123 | detective_regions = user_regions.split(',') 124 | logging.info(f'Enabling members in these regions: {detective_regions}') 125 | else: 126 | detective_regions = session.get_available_regions('detective') 127 | logging.info( 128 | f'Enabling members in all available Detective regions {detective_regions}') 129 | return detective_regions 130 | 131 | 132 | def assume_role(aws_account_number: str, role_name: str) -> boto3.Session: 133 | """ 134 | Assumes the provided role in each account and returns a Detective client. 135 | 136 | Args: 137 | - aws_account_number: AWS Account Number 138 | - role_name: Role to assume in target account 139 | 140 | Returns: 141 | Detective client in the specified AWS Account and Region 142 | """ 143 | try: 144 | # Beginning the assume role process for account 145 | sts_client = boto3.client('sts') 146 | 147 | # Get the current partition 148 | partition = sts_client.get_caller_identity()['Arn'].split(":")[1] 149 | 150 | response = sts_client.assume_role( 151 | RoleArn='arn:{}:iam::{}:role/{}'.format( 152 | partition, 153 | aws_account_number, 154 | role_name 155 | ), 156 | RoleSessionName='EnableDetective' 157 | ) 158 | # Storing STS credentials 159 | session = boto3.Session( 160 | aws_access_key_id=response['Credentials']['AccessKeyId'], 161 | aws_secret_access_key=response['Credentials']['SecretAccessKey'], 162 | aws_session_token=response['Credentials']['SessionToken'] 163 | ) 164 | except Exception as e: 165 | logging.exception(f'exception: {e}') 166 | 167 | logging.info(f"Assumed session for {aws_account_number}.") 168 | 169 | return session 170 | 171 | 172 | def get_graphs(d_client: botocore.client.BaseClient) -> typing.List[str]: 173 | """ 174 | Get graphs in a specified region. 175 | 176 | Args: 177 | - d_client: Detective boto3 client generated from the master session. 178 | 179 | Returns: 180 | List of graph Arns. 181 | """ 182 | try: 183 | response = d_client.list_graphs() 184 | except botocore.exceptions.EndpointConnectionError: 185 | logging.exception(f'exception: {e}') 186 | return [] 187 | 188 | # use .get function to avoid KeyErrors when a dictionary key doesn't exist 189 | # it returns an empty list instead. 190 | # map iterates over all elements under a list and applied a function to them, 191 | # in this specific case, the element 'Arn' is extracted from the dictionary 192 | # (graphlist is a list of dictionaries) 193 | return [x['Arn'] for x in response.get('GraphList', [])] 194 | 195 | 196 | def get_members(d_client: botocore.client.BaseClient, graphs: typing.List[str]) ->\ 197 | (typing.Dict[str, typing.Set[str]], typing.Dict[str, typing.Set[str]]): 198 | """ 199 | Get member accounts for all behaviour graphs in a region. 200 | 201 | Args: 202 | - d_client: Detective boto3 client generated from the master session. 203 | - graphs: List of graphs arns 204 | 205 | Returns: 206 | Two dictionaries: one with all account ids, other with the ones pending to accept 207 | the invitation. 208 | """ 209 | try: 210 | # itertools.tee creates two independent iterators from a single one. This way 211 | # we can iterate the iterator twice: one to return all elements and other to return 212 | # the ones pending to be invited. 213 | #### 214 | # check the value of NextToken in the response. if it is non-null, pass it back into a subsequent list_members call (and keep doing this until a null token is returned) 215 | def _master_memberList(g: str) -> typing.List[typing.Dict]: 216 | # create a list to append the member accounts 217 | memberAccounts = [] 218 | # create a dictionary for the nextToken from each call 219 | tokenTracker = {} 220 | # loop through list_members call results and take action for each returned result 221 | while True: 222 | # list_members of graph "g" and return the first 100 results 223 | members = d_client.list_members( 224 | GraphArn=g, MaxResults=100, **tokenTracker) 225 | # add the returned list members to the list 226 | memberAccounts.extend(members['MemberDetails']) 227 | # if the returned results have a "NextToken" key then use it to query again 228 | if 'NextToken' in members: 229 | tokenTracker['NextToken'] = members['NextToken'] 230 | # if the returned results do not have a "NextToken" key then exit the loop 231 | else: 232 | break 233 | # return members list. 234 | # The return statement doesn't need () 235 | return memberAccounts 236 | except Exception as e: 237 | logging.exception(f'exception when getting memebers: {e}') 238 | 239 | # iterate through each list and return results 240 | all_ac, pending = itertools.tee((g, _master_memberList(g)) 241 | for g in graphs) 242 | 243 | return ({g: {x['AccountId'] for x in v} for g, v in all_ac}, 244 | {g: {x['AccountId'] for x in v if x['Status'] == 'INVITED'} for g, v in pending}) 245 | 246 | 247 | def create_members(d_client: botocore.client.BaseClient, graph_arn: str, disable_email: bool, account_ids: typing.Set[str], 248 | account_csv: typing.Dict[str, str]) -> typing.Set[str]: 249 | """ 250 | Creates member accounts for all accounts in the csv that are not present in the graph member set. 251 | 252 | Args: 253 | - d_client: Detective boto3 client generated from the master session. 254 | - graph_arn: Graph to add members to. 255 | - account_ids: Already present account ids in the graph. 256 | - account_csv: Accounts read from the CSV input file. 257 | 258 | Returns: 259 | Set with the IDs of the successfully created accounts. 260 | """ 261 | try: 262 | # I'm calculating set difference: the elements that are present in the CSV and that are not 263 | # present in the account_ids set. 264 | set_difference = account_csv.keys() - account_ids 265 | if not set_difference: 266 | logging.info(f'No new members to create in graph {graph_arn}.') 267 | return set() 268 | 269 | logging.info(f'Creating member accounts in graph {graph_arn} ' 270 | f'for accounts {", ".join(set_difference)}.') 271 | 272 | new_members = [{'AccountId': x, 'EmailAddress': account_csv[x]} 273 | for x in set_difference] 274 | response = d_client.create_members(GraphArn=graph_arn, 275 | Message='Automatically generated invitation', 276 | Accounts=new_members, 277 | DisableEmailNotification=disable_email) 278 | for error in response['UnprocessedAccounts']: 279 | logging.exception(f'Could not create member for account {error["AccountId"]} in ' 280 | f'graph {graph_arn}: {error["Reason"]}') 281 | except Exception as e: 282 | logging.exception(f'exception when getting memebers: {e}') 283 | return {x['AccountId'] for x in response['Members']} 284 | 285 | 286 | def accept_invitations(role: str, accounts: typing.Set[str], graph: str, region: str) -> typing.NoReturn: 287 | """ 288 | Accept invitation for a list of accounts in a given graph. 289 | 290 | Args: 291 | - role: Role to assume when accepting the invitation. 292 | - accounts: Set of accounts pending to accept. 293 | - graph: Graph the accounts are being invited to. 294 | - region: Region for the client 295 | """ 296 | try: 297 | for account in accounts: 298 | logging.info( 299 | f'Accepting invitation for account {account} in graph {graph}.') 300 | session = assume_role(account, role) 301 | local_client = session.client('detective', region_name=region) 302 | local_client.accept_invitation(GraphArn=graph) 303 | except Exception as e: 304 | logging.exception(f'error accepting invitation {e.args}') 305 | 306 | def enable_detective(d_client: botocore.client.BaseClient, region: str, tags: dict = {}): 307 | graphs = get_graphs(d_client) 308 | 309 | if not graphs: 310 | confirm = input('Should Amazon Detective be enabled in {}? Enter [Y/N]: '.format(region)) 311 | 312 | if confirm == 'Y' or confirm == 'y': 313 | logging.info(f'Enabling Amazon Detective in {region}' + (f' with tags {tags}' if tags else '')) 314 | graphs = [d_client.create_graph(Tags=tags)['GraphArn']] 315 | else: 316 | logging.info(f'Skipping {region}') 317 | return None 318 | logging.info(f'Amazon Detective is enabled in region {region}') 319 | 320 | return graphs 321 | 322 | def chunked(it, size): 323 | it = iter(it) 324 | while True: 325 | p = tuple(itertools.islice(it, size)) 326 | if not p: 327 | break 328 | yield p 329 | 330 | if __name__ == '__main__': 331 | args = setup_command_line() 332 | aws_account_dict = read_accounts_csv(args.input_file) 333 | 334 | try: 335 | session = boto3.session.Session() 336 | detective_regions = get_regions(session, args.enabled_regions) 337 | master_session = assume_role(args.master_account, args.assume_role) 338 | except NameError as e: 339 | # logging.exception prints the full traceback, logging.error just prints the error message. 340 | # In this case, the NameError is handled in a specific way: it is an expected exception 341 | # and the traceback doesn't add any value to the error message. 342 | logging.error(f'Master account is not defined: {e.args}') 343 | except Exception as e: 344 | # in this case, there has been an unhandled exception, something that wasn't estimated. 345 | # Having the error traceback helps us know what happened and help us find a solution 346 | # for the bug. The code should never arrive to this except clause, but if it does 347 | # we want as much information as possible and that's why we use logging.traceback. 348 | # In this case the traceback adds LOTS of value. 349 | logging.exception(f'error creating session {e.args}') 350 | 351 | #Chunk the list of accounts in the .csv into batches of 50 due to the API limitation of 50 accounts per invokation 352 | for chunk_tuple in chunked(aws_account_dict.items(), 50): 353 | chunk = {x: y for x, y in chunk_tuple} 354 | 355 | for region in detective_regions: 356 | try: 357 | d_client = master_session.client('detective', region_name=region) 358 | graphs = enable_detective(d_client, region, args.tags) 359 | 360 | if graphs is None: 361 | continue 362 | 363 | try: 364 | all_members, pending = get_members(d_client, graphs) 365 | 366 | for graph, members in all_members.items(): 367 | new_accounts = create_members( 368 | d_client, graph, args.disable_email, members, chunk) 369 | print("Sleeping for 5s to allow new members' invitations to propagate.") 370 | time.sleep(5) 371 | accept_invitations(args.assume_role, itertools.chain( 372 | new_accounts, pending[graph]), graph, region) 373 | except NameError as e: 374 | logging.error(f'account is not defined: {e}') 375 | except Exception as e: 376 | logging.exception(f'unable to accept invitiation: {e}') 377 | 378 | except NameError as e: 379 | logging.error(f'account is not defined: {e}') 380 | except Exception as e: 381 | logging.exception(f'error with region {region}: {e}') 382 | --------------------------------------------------------------------------------