diff --git a/repokid/__init__.py b/repokid/__init__.py index 453fdf6fe..fa9421be1 100644 --- a/repokid/__init__.py +++ b/repokid/__init__.py @@ -20,7 +20,7 @@ import import_string -__version__ = "0.14.2" +__version__ = "0.14.3" def init_config(): diff --git a/repokid/exceptions/__init__.py b/repokid/exceptions/__init__.py index 150a417d0..a3e59b85b 100644 --- a/repokid/exceptions/__init__.py +++ b/repokid/exceptions/__init__.py @@ -1,2 +1,10 @@ class UnexpectedDynamoUpdateValue(Exception): pass + + +class BlocklistError(Exception): + pass + + +class AardvarkError(Exception): + pass diff --git a/repokid/filters/blocklist/__init__.py b/repokid/filters/blocklist/__init__.py index f3b160883..03347b293 100644 --- a/repokid/filters/blocklist/__init__.py +++ b/repokid/filters/blocklist/__init__.py @@ -1,9 +1,9 @@ import json import logging -import sys import botocore from cloudaux.aws.sts import boto3_cached_conn +from repokid.exceptions import BlocklistError from repokid.filters import Filter LOGGER = logging.getLogger("repokid") @@ -27,30 +27,26 @@ def get_blocklist_from_bucket(bucket_config): blocklist_json = json.loads(blocklist) # Blocklist problems are really bad and we should quit rather than silently continue except (botocore.exceptions.ClientError, AttributeError): - LOGGER.error( + LOGGER.critical( "S3 blocklist config was set but unable to connect retrieve object, quitting" ) - sys.exit(1) + raise BlocklistError("Could not retrieve blocklist") except ValueError: - LOGGER.error( + LOGGER.critical( "S3 blocklist config was set but the returned file is bad, quitting" ) - sys.exit(1) - if set(blocklist_json.keys()) != set(["arns", "names"]): - LOGGER.error("S3 blocklist file is malformed, quitting") - sys.exit(1) + raise BlocklistError("Could not parse blocklist") + if set(blocklist_json.keys()) != {"arns", "names"}: + LOGGER.critical("S3 blocklist file is malformed, quitting") + raise BlocklistError("Could not parse blocklist") return blocklist_json class BlocklistFilter(Filter): - def __init__(self, config=None): - blocklist_json = None - bucket_config = config.get( - "blocklist_bucket", config.get("blacklist_bucket", None) - ) - if bucket_config: - blocklist_json = get_blocklist_from_bucket(bucket_config) + blocklist_json = None + def __init__(self, config=None): + super().__init__(config) current_account = config.get("current_account") or None if not current_account: LOGGER.error("Unable to get current account for Blocklist Filter") @@ -63,20 +59,33 @@ def __init__(self, config=None): [rolename.lower() for rolename in config.get("all", [])] ) - if blocklist_json: + if BlocklistFilter.blocklist_json: blocklisted_role_names.update( [ name.lower() - for name, accounts in blocklist_json["names"].items() + for name, accounts in BlocklistFilter.blocklist_json[ + "names" + ].items() if ("all" in accounts or config.get("current_account") in accounts) ] ) self.blocklisted_arns = ( - set() if not blocklist_json else blocklist_json.get("arns", []) + set() + if not BlocklistFilter.blocklist_json + else BlocklistFilter.blocklist_json.get("arns", []) ) self.blocklisted_role_names = blocklisted_role_names + @classmethod + def init_blocklist(cls, config): + if not cls.blocklist_json: + bucket_config = config.get( + "blocklist_bucket", config.get("blacklist_bucket", None) + ) + if bucket_config: + cls.blocklist_json = get_blocklist_from_bucket(bucket_config) + def apply(self, input_list): blocklisted_roles = [] diff --git a/repokid/utils/aardvark.py b/repokid/utils/aardvark.py index c2047bc7e..b47042298 100644 --- a/repokid/utils/aardvark.py +++ b/repokid/utils/aardvark.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import logging -import sys +from repokid.exceptions import AardvarkError import requests LOGGER = logging.getLogger("repokid") @@ -52,12 +52,12 @@ def get_aardvark_data(aardvark_api_location, account_number=None, arn=None): aardvark_api_location, params=params, json=payload ) except requests.exceptions.RequestException as e: - LOGGER.error("Unable to get Aardvark data: {}".format(e)) - sys.exit(1) + LOGGER.exception("Unable to get Aardvark data: {}".format(e)) + raise AardvarkError("Unable to get aardvark data") else: if r_aardvark.status_code != 200: - LOGGER.error("Unable to get Aardvark data") - sys.exit(1) + LOGGER.exception("Unable to get Aardvark data") + raise AardvarkError("Unable to get aardvark data") response_data.update(r_aardvark.json()) # don't want these in our Aardvark data diff --git a/repokid/utils/dynamo.py b/repokid/utils/dynamo.py index 0fd1397e4..149878928 100644 --- a/repokid/utils/dynamo.py +++ b/repokid/utils/dynamo.py @@ -1,7 +1,7 @@ import copy import datetime +from functools import wraps import logging -import sys import boto3 from botocore.exceptions import ClientError as BotoClientError @@ -13,12 +13,12 @@ def catch_boto_error(func): + @wraps(func) def decorated_func(*args, **kwargs): try: return func(*args, **kwargs) except BotoClientError as e: - LOGGER.error("Dynamo table error: {}".format(e)) - sys.exit(1) + LOGGER.exception("Dynamo table error: {}".format(e)) return decorated_func @@ -234,16 +234,18 @@ def set_role_data(dynamo_table, role_id, update_keys): value = _empty_string_to_dynamo_replace(value) - update_expression += "#expr{} = :val{}".format(count, count) - expression_attribute_names["#expr{}".format(count)] = key - expression_attribute_values[":val{}".format(count)] = value + update_expression += f"#expr{count} = :val{count}" + expression_attribute_names[f"#expr{count}"] = key + expression_attribute_values[f":val{count}"] = value - dynamo_table.update_item( - Key={"RoleId": role_id}, - UpdateExpression=update_expression, - ExpressionAttributeNames=expression_attribute_names, - ExpressionAttributeValues=expression_attribute_values, - ) + update_item_inputs = { + "Key": {"RoleId": role_id}, + "UpdateExpression": update_expression, + "ExpressionAttributeNames": expression_attribute_names, + "ExpressionAttributeValues": expression_attribute_values, + } + LOGGER.debug("updating dynamodb with inputs %s", update_item_inputs) + dynamo_table.update_item(**update_item_inputs) def store_initial_role_data(