Don't exit on errors, only fetch blocklist once (#259)

repokid/__init__.py

@@ -20,7 +20,7 @@
 
 import import_string
 
-__version__ = "0.14.2"
+__version__ = "0.14.3"
 
 
 def init_config():

repokid/exceptions/__init__.py

@@ -1,2 +1,10 @@
 class UnexpectedDynamoUpdateValue(Exception):
     pass
+
+
+class BlocklistError(Exception):
+    pass
+
+
+class AardvarkError(Exception):
+    pass

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 = []
 

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

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(