A few enhancements

- Create a new function that can update individual role data. This is useful so we can get an accurate count of repoable, etc permissions after a repo or rollback. - Find roles by name using a secondary index now (should improve performance significantly) - Change description when a role is repoed to include the date - When showing warnings about a permission being on the no repo list, show info instead and show the role name - Only update stats when the count or repoable services are different than before - Fix a bug where display role fails if a bad role name is given - Fix a bug in the OptOut filter
Netflix · Aug 28, 2017 · e5cf04a · e5cf04a
1 parent de09d19
commit e5cf04a
Show file tree

Hide file tree

Showing 7 changed files with 168 additions and 76 deletions.
diff --git a/README.md b/README.md
@@ -66,7 +66,8 @@ RepokidRole:
         "iam:ListInstanceProfilesForRole",
         "iam:ListRolePolicies",
         "iam:ListRoles",
-        "iam:PutRolePolicy"
+        "iam:PutRolePolicy",
+        "iam:UpdateRoleDescription"
       ],
       "Effect": "Allow",
       "Resource": "*"

diff --git a/repokid/__init__.py b/repokid/__init__.py
@@ -15,7 +15,7 @@
 import logging.config
 import os
 
-__version__ = '0.7.1'
+__version__ = '0.7.2'
 
 
 def init_config():

diff --git a/repokid/cli/repokid_cli.py b/repokid/cli/repokid_cli.py
@@ -33,11 +33,13 @@
 import datetime
 import json
 import pprint
+import re
 import requests
 import sys
 
 import botocore
 from cloudaux.aws.iam import list_roles, get_role_inline_policies
+from cloudaux.aws.sts import sts_conn
 from docopt import docopt
 import import_string
 from tabulate import tabulate
@@ -163,15 +165,17 @@ def _generate_default_config(filename=None):
     return config_dict
 
 
-def _get_aardvark_data(account_number, aardvark_api_location):
+def _get_aardvark_data(aardvark_api_location, account_number=None, arn=None):
     """
-    Make a request to the Aardvark server to get all data about a given account.
+    Make a request to the Aardvark server to get all data about a given account or ARN.
     We'll request in groups of PAGE_SIZE and check the current count to see if we're done. Keep requesting as long as
     the total count (reported by the API) is greater than the number of pages we've received times the page size.  As
     we go, keeping building the dict and return it when done.
 
     Args:
+        aardvark_api_location
         account_number (string): Used to form the phrase query for Aardvark so we only get data for the account we want
+        arn (string)
 
     Returns:
         dict: Aardvark data is a dict with the role ARN as the key and a list of services as value
@@ -181,7 +185,12 @@ def _get_aardvark_data(account_number, aardvark_api_location):
     PAGE_SIZE = 1000
     page_num = 1
 
-    payload = {'phrase': '{}'.format(account_number)}
+    if account_number:
+        payload = {'phrase': '{}'.format(account_number)}
+    elif arn:
+        payload = {'arn': [arn]}
+    else:
+        return
     while True:
         params = {'count': PAGE_SIZE, 'page': page_num}
         try:
@@ -206,6 +215,72 @@ def _get_aardvark_data(account_number, aardvark_api_location):
     return response_data
 
 
+@sts_conn('iam')
+def _update_repoed_description(role_name, client=None):
+    description = None
+    try:
+        description = client.get_role(RoleName=role_name)['Role'].get('Description', '')
+    except KeyError:
+        return
+    date_string = datetime.datetime.utcnow().strftime('%m/%d/%y')
+    if '; Repokid repoed' in description:
+        new_description = re.sub(r'; Repokid repoed [0-9]{2}\/[0-9]{2}\/[0-9]{2}', '; Repokid repoed {}'.format(
+                                 date_string), description)
+    else:
+        new_description = description + ' ; Repokid repoed {}'.format(date_string)
+    # IAM role descriptions have a max length of 1000, if our new length would be longer, skip this
+    if len(new_description < 1000):
+        client.update_role_description(RoleName=role_name, Description=new_description)
+    else:
+        LOGGER.erorr('Unable to set repo description ({}) for role {}, length would be too long'.format(
+            new_description, role_name))
+
+
+def _update_role_data(role, dynamo_table, account_number, config, conn, source, add_no_repo=True):
+    """
+    Perform a scaled down version of role update, this is used to get an accurate count of repoable permissions after
+    a rollback or repo.
+
+    Does update:
+     - Policies
+     - Aardvark data
+     - Total permissions
+     - Repoable permissions
+     - Repoable services
+     - Stats
+
+    Does not update:
+     - Filters
+     - Active/inactive roles
+
+    Args:
+        role (Role)
+        dynamo_table
+        account_number
+        conn (dict)
+        source: repo, rollback, etc
+        add_no_repo: if set to True newly discovered permissions will be added to no repo list
+
+    Returns:
+        None
+    """
+    current_policies = get_role_inline_policies(role.as_dict(), **conn) or {}
+    roledata.update_role_data(dynamo_table, account_number, role, current_policies, source=source,
+                              add_no_repo=add_no_repo)
+    aardvark_data = _get_aardvark_data(config['aardvark_api_location'], arn=role.arn)
+
+    if not aardvark_data:
+        return
+
+    role.aa_data = aardvark_data[role.arn]
+    roledata._calculate_repo_scores([role], config['filter_config']['AgeFilter']['minimum_age'])
+    set_role_data(dynamo_table, role.role_id, {'AAData': role.aa_data,
+                                               'TotalPermissions': role.total_permissions,
+                                               'RepoablePermissions': role.repoable_permissions,
+                                               'RepoableServices': role.repoable_services})
+    roledata.update_stats(dynamo_table, [role], source=source)
+
+
 # inspiration from https://github.com/slackhq/python-rtmbot/blob/master/rtmbot/core.py
 class FilterPlugins(object):
     """
@@ -303,7 +378,7 @@ def update_role_cache(account_number, dynamo_table, config):
         set_role_data(dynamo_table, role.role_id, {'DisqualifiedBy': role.disqualified_by})
 
     LOGGER.info('Getting data from Aardvark')
-    aardvark_data = _get_aardvark_data(account_number, config['aardvark_api_location'])
+    aardvark_data = _get_aardvark_data(config['aardvark_api_location'], account_number=account_number)
 
     LOGGER.info('Updating with Aardvark data')
     for role in roles:
@@ -403,13 +478,11 @@ def display_role(account_number, role_name, dynamo_table, config):
         None
     """
     role_id = find_role_in_cache(dynamo_table, account_number, role_name)
-    role_data = get_role_data(dynamo_table, role_id)
-
-    if not role_data:
+    if not role_id:
         LOGGER.warn('Could not find role with name {}'.format(role_name))
         return
-    else:
-        role = Role(role_data)
+
+    role = Role(get_role_data(dynamo_table, role_id))
 
     print "\n\nRole repo data:"
     headers = ['Name', 'Refreshed', 'Disqualified By', 'Can be repoed', 'Permissions', 'Repoable', 'Repoed', 'Services']
@@ -454,7 +527,8 @@ def display_role(account_number, role_name, dynamo_table, config):
 
     permissions = roledata._get_role_permissions(role, warn_unknown_perms=warn_unknown_permissions)
     if len(role.disqualified_by) == 0:
-        repoable_permissions = roledata._get_repoable_permissions(permissions, role.aa_data, role.no_repo_permissions,
+        repoable_permissions = roledata._get_repoable_permissions(role.role_name, permissions, role.aa_data,
+                                                                  role.no_repo_permissions,
                                                                   config['filter_config']['AgeFilter']['minimum_age'])
 
     print "Repoable services:"
@@ -535,7 +609,8 @@ def repo_role(account_number, role_name, dynamo_table, config, commit=False):
         return
 
     permissions = roledata._get_role_permissions(role)
-    repoable_permissions = roledata._get_repoable_permissions(permissions, role.aa_data, role.no_repo_permissions,
+    repoable_permissions = roledata._get_repoable_permissions(role.role_name, permissions, role.aa_data,
+                                                              role.no_repo_permissions,
                                                               config['filter_config']['AgeFilter']['minimum_age'])
     repoed_policies, deleted_policy_names = roledata._get_repoed_policy(role.policies[-1]['Policy'],
                                                                         repoable_permissions)
@@ -589,18 +664,8 @@ def repo_role(account_number, role_name, dynamo_table, config, commit=False):
 
     if not errors:
         set_role_data(dynamo_table, role.role_id, {'Repoed': datetime.datetime.utcnow().isoformat()})
-
-        # update total and repoable permissions and services
-        role.total_permissions = len(roledata._get_role_permissions(role))
-        role.repoable_permissions = 0
-        role.repoable_services = []
-        set_role_data(dynamo_table, role.role_id, {'TotalPermissions': role.total_permissions,
-                                                   'RepoablePermissions': 0,
-                                                   'RepoableServices': []})
-
-        # update stats
-        roledata.update_stats(dynamo_table, [role], source='Repo')
-
+        _update_repoed_description(role.role_name, **conn)
+        _update_role_data(role, dynamo_table, account_number, config, conn, source='Repo', add_no_repo=False)
         LOGGER.info('Successfully repoed role: {}'.format(role.role_name))
     return errors
 
@@ -689,14 +754,7 @@ def rollback_role(account_number, role_name, dynamo_table, config, selection=Non
                 LOGGER.error(message)
                 errors.append(message)
 
-    # TODO: possibly update the total and repoable permissions here, we'd have to get Aardvark and all that
-
-    current_policies = get_role_inline_policies(role.as_dict(), **conn) or {}
-    roledata.add_new_policy_version(dynamo_table, role, current_policies, 'Restore')
-    role.total_permissions = len(roledata._get_role_permissions(role))
-
-    # update stats
-    roledata.update_stats(dynamo_table, [role], source='Restore')
+    _update_role_data(role, dynamo_table, account_number, config, conn, source='Restore', add_no_repo=False)
 
     if not errors:
         LOGGER.info('Successfully restored selected version of role policies')

diff --git a/repokid/filters/optout/__init__.py b/repokid/filters/optout/__init__.py
@@ -1,4 +1,4 @@
-import time.time
+import time
 
 from repokid.cli.repokid_cli import Filter
 

diff --git a/repokid/tests/test_roledata.py b/repokid/tests/test_roledata.py
@@ -140,7 +140,7 @@ def test_get_repoable_permissions(self):
 
         no_repo_permissions = {'service_4:action_1': time.time()-1, 'service_4:action_2': time.time()+1000}
 
-        repoable_permissions = repokid.utils.roledata._get_repoable_permissions(permissions, aa_data,
+        repoable_permissions = repokid.utils.roledata._get_repoable_permissions('test_name', permissions, aa_data,
                                                                                 no_repo_permissions, minimum_age)
         # service_1:action_3 and action_4 are unsupported actions, service_2 is an unsupported service, service_3
         # was used too recently, service_4 action 2 is in no_repo_permissions and not expired

diff --git a/repokid/utils/dynamo.py b/repokid/utils/dynamo.py
@@ -69,38 +69,54 @@ def dynamo_get_or_create_table(**dynamo_config):
                 'KeyType': 'HASH'  # Partition key
             }],
             AttributeDefinitions=[{
-                'AttributeName': 'RoleId',
-                'AttributeType': 'S'
-            }],
+                    'AttributeName': 'RoleId',
+                    'AttributeType': 'S'
+                },
+                {
+                    'AttributeName': 'RoleName',
+                    'AttributeType': 'S'
+                },
+                {
+                    'AttributeName': 'Account',
+                    'AttributeType': 'S'
+                }],
             ProvisionedThroughput={
-                'ReadCapacityUnits': 5,
-                'WriteCapacityUnits': 5
-            })
-
-        table.meta.client.get_waiter('table_exists').wait(TableName='repokid_roles')
-
-        # need a global secondary index to list all role IDs for a given account number
-        table.update(
-            AttributeDefinitions=[{
-                'AttributeName': 'Account',
-                'AttributeType': 'S'
-             }],
-            GlobalSecondaryIndexUpdates=[{
-                'Create': {
+                'ReadCapacityUnits': 50,
+                'WriteCapacityUnits': 50
+            },
+            GlobalSecondaryIndexes=[
+                {
                     'IndexName': 'Account',
-                    'KeySchema': [{
-                        'AttributeName': 'Account',
-                        'KeyType': 'HASH'
-                    }],
+                    'KeySchema': [
+                        {
+                            'AttributeName': 'Account',
+                            'KeyType': 'HASH'
+                        }
+                    ],
                     'Projection': {
-                        'NonKeyAttributes': ['RoleId'],
-                        'ProjectionType': 'INCLUDE'
+                        'ProjectionType': 'KEYS_ONLY',
                     },
                     'ProvisionedThroughput': {
-                        'ReadCapacityUnits': 2,
-                        'WriteCapacityUnits': 2
+                        'ReadCapacityUnits': 10,
+                        'WriteCapacityUnits': 10
                     }
-                }}])
+                },
+                {
+                    'IndexName': 'RoleName',
+                    'KeySchema': [
+                        {
+                            'AttributeName': 'RoleName',
+                            'KeyType': 'HASH'
+                        }
+                    ],
+                    'Projection': {
+                        'ProjectionType': 'KEYS_ONLY',
+                    },
+                    'ProvisionedThroughput': {
+                        'ReadCapacityUnits': 10,
+                        'WriteCapacityUnits': 10
+                    }
+                }])
 
     except BotoClientError as e:
         if "ResourceInUseException" in e.message:
@@ -120,10 +136,20 @@ def find_role_in_cache(dynamo_table, account_number, role_name):
     Returns:
         string: RoleID for active role with name in given account, else None
     """
-    for roleID in role_ids_for_account(dynamo_table, account_number):
-        role_data = get_role_data(dynamo_table, roleID, fields=['RoleName', 'Active', 'RoleId'])
-        if role_data['RoleName'].lower() == role_name.lower() and role_data['Active']:
-            return role_data['RoleId']
+    results = dynamo_table.query(IndexName='RoleName',
+                                 KeyConditionExpression='RoleName = :rn',
+                                 ExpressionAttributeValues={':rn': role_name})
+    role_id_candidates = [return_dict['RoleId'] for return_dict in results.get('Items')]
+
+    if len(role_id_candidates) > 1:
+        for role_id in role_id_candidates:
+            role_data = get_role_data(dynamo_table, role_id, fields=['Active'])
+            if role_data['Active']:
+                return role_id
+    elif len(role_id_candidates) == 1:
+        return role_id_candidates[0]
+    else:
+        return None
 
 
 @catch_boto_error
@@ -160,14 +186,12 @@ def role_ids_for_account(dynamo_table, account_number):
     role_ids = set()
 
     results = dynamo_table.query(IndexName='Account',
-                                 ProjectionExpression='RoleId',
                                  KeyConditionExpression='Account = :act',
                                  ExpressionAttributeValues={':act': account_number})
     role_ids.update([return_dict['RoleId'] for return_dict in results.get('Items')])
 
     while 'LastEvaluatedKey' in results:
         results = dynamo_table.query(IndexName='Account',
-                                     ProjectionExpression='RoleId',
                                      KeyConditionExpression='Account = :act',
                                      ExpressionAttributeValues={':act': account_number},
                                      ExclusiveStartKey=results.get('LastEvaluatedKey'))