diff --git a/automation_rules/list-automation-rules.py b/automation_rules/list-automation-rules.py index dc5330a..8d025cb 100755 --- a/automation_rules/list-automation-rules.py +++ b/automation_rules/list-automation-rules.py @@ -15,37 +15,89 @@ OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ - -import boto3 import argparse import json +import boto3 +from botocore.config import Config from botocore.exceptions import ClientError -parser = argparse.ArgumentParser(description='List deployed automation rules in one or many regions.') -parser.add_argument('--deployed_regions', type=str, required=False, help="Comma separated list of regions to list rules from. If not specified, rules from all regions will be retrieved.") +# Custom Config for adaptive retry strategy +custom_config = Config(retries={"max_attempts": 10, "mode": "adaptive"}) + +# Argument Parsing +parser = argparse.ArgumentParser( + description="List and save deployed automation rules in one or many regions." +) +parser.add_argument( + "--deployed_regions", + type=str, + required=False, + help="Comma-separated list of regions to list rules from. If not specified, rules from all regions will be retrieved.", +) +parser.add_argument( + "--save-rules-json", + action="store_true", + help="Save the rules as JSON files in the current directory with region name appended.", +) args = parser.parse_args() -# Getting SecurityHub regions + +# Function to save rule to JSON +def save_rule_to_json(rule_name, rule_data, region): + """ + Saves the given rule data to a JSON file with the given name and region. + + Args: + rule_name (str): The name of the rule. + rule_data (dict): The data for the rule to be saved. + region (str): The region where the rule is saved. + + Returns: + None + """ + filename = f"{rule_name}_{region}.json" + with open(filename, "w", encoding="UTF-8") as file: + json.dump(rule_data, file, indent=4, sort_keys=True, default=str) + print(f"Saved rule '{rule_name}' to JSON file in {region}.") + + +# Boto3 Session session = boto3.session.Session() -securityhub_regions = [] -if args.deployed_regions: - securityhub_regions = [str(item) for item in args.deployed_regions.split(",")] - print("Listing rules in these regions: {}".format(securityhub_regions)) -else: - securityhub_regions = session.get_available_regions('securityhub') - print("Listing rules in all available SecurityHub regions {}".format(securityhub_regions)) +# Determine SecurityHub regions +securityhub_regions = ( + args.deployed_regions.split(",") + if args.deployed_regions + else session.get_available_regions("securityhub") +) +print("Processing these regions: {}".format(securityhub_regions)) +# Process each region for aws_region in securityhub_regions: print("*******************************************") print("Retrieving rules from region: ", aws_region) - sh_client = session.client("securityhub", region_name=aws_region) + sh_client = session.client( + "securityhub", region_name=aws_region, config=custom_config + ) try: - rules = sh_client.list_automation_rules(MaxResults=100).get( - "AutomationRulesMetadata", [] - ) + rule_arns = [] + next_token = "" + while True: + if next_token: + response = sh_client.list_automation_rules( + MaxResults=100, NextToken=next_token + ) + else: + response = sh_client.list_automation_rules(MaxResults=100) + + rules = response.get("AutomationRulesMetadata", []) + rule_arns.extend([rule["RuleArn"] for rule in rules]) + + next_token = response.get("NextToken") + if not next_token: + break if rules: for msg in rules: rule_arn = msg["RuleArn"] @@ -59,7 +111,17 @@ print("Rule Order: ", rule_order) else: print("No rules in this region") + # Getting full rule definitions + if rule_arns: + rule_defs_response = sh_client.batch_get_automation_rules( + AutomationRulesArns=rule_arns + ) + rule_defs = rule_defs_response.get("Rules", []) # Corrected key here + for rule_def in rule_defs: + rule_name = rule_def["RuleName"] + print(f"Processing rule: {rule_name}") + if args.save_rules_json: + save_rule_to_json(rule_name, rule_def, aws_region) except ClientError as e: - print("Error Processing Region {}".format(aws_region)) - print(e) + print(f"Error processing region {aws_region}: {e}")