Refactoring to use the ARN class.

security_monkey/auditors/iam/iam_role.py

@@ -26,8 +26,6 @@
 from security_monkey.watchers.iam.managed_policy import ManagedPolicy
 from security_monkey.datastore import Account
 
-import re
-
 
 class IAMRoleAuditor(IAMPolicyAuditor):
     index = IAMRole.index
@@ -58,19 +56,18 @@ def check_statement(statement):
                         if aws and aws == "*":
                             self.add_issue(10, tag, iamrole_item,
                                            notes=json.dumps(statement))
-                    elif aws and type(aws) is list:
-                        for entry in aws:
-                            if entry == "*":
-                                self.add_issue(10, tag, iamrole_item,
-                                               notes=json.dumps(statement))
+                        elif aws and type(aws) is list:
+                            for entry in aws:
+                                if entry == "*":
+                                    self.add_issue(10, tag, iamrole_item,
+                                                   notes=json.dumps(statement))
 
         assume_role_policy = iamrole_item.config.get("AssumeRolePolicyDocument", {})
         statement = assume_role_policy.get("Statement", [])
-        if type(statement) is list:
-            for single_statement in statement:
-                check_statement(single_statement)
-        elif type(statement) is dict:
-            check_statement(statement)
+        if type(statement) is dict:
+            statement = [statement]
+        for single_statement in statement:
+            check_statement(single_statement)
 
     def check_assume_role_from_unknown_account(self, iamrole_item):
         """
@@ -79,13 +76,19 @@ def check_assume_role_from_unknown_account(self, iamrole_item):
 
         def check_statement(statement):
 
-            def check_account_in_arn(arn):
-                m = re.match(r'arn:aws:iam::([0-9*]+):', arn)
-                if m.group(1):
-                    account = Account.query.filter(Account.number == m.group(1)).first()
+            def check_account_in_arn(input):
+                from security_monkey.common.arn import ARN
+                arn = ARN(input)
+
+                if arn.error:
+                    print('Could not parse ARN in Trust Policy: {arn}'.format(arn=input))
+
+                if not arn.error and arn.account_number:
+                    account = Account.query.filter(Account.number == arn.account_number).first()
                     if not account:
-                        tag = "{0} allows assume-role from an Unknown Account ({1})".format(self.i_am_singular,
-                                                                                           m.group(1))
+                        tag = "IAM Role allows assume-role from an " \
+                            + "Unknown Account ({account_number})".format(
+                            account_number=arn.account_number)
                         self.add_issue(10, tag, iamrole_item, notes=json.dumps(statement))
 
             action = statement.get("Action", None)
@@ -105,11 +108,10 @@ def check_account_in_arn(arn):
 
         assume_role_policy = iamrole_item.config.get("AssumeRolePolicyDocument", {})
         statement = assume_role_policy.get("Statement", [])
-        if type(statement) is list:
-            for single_statement in statement:
-                check_statement(single_statement)
-        elif type(statement) is dict:
-            check_statement(statement)
+        if type(statement) is dict:
+            statement = [statement]
+        for single_statement in statement:
+            check_statement(single_statement)
 
     def check_star_privileges(self, iamrole_item):
         """