Add basic filtering by role-arn (#3)

CHANGELOG.md

@@ -4,6 +4,12 @@ This changelog contains a loose collection of changes in every release including
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 
+## 0.2.0
+
+### Added
+
+* Basic filtering for role-arns when generating policy (#3)
+
 ## 0.1.0
 
 _Initial Release_

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.1.0
+current_version = 0.2.0
 commit = True
 tag = True
 

setup.py

@@ -26,7 +26,7 @@
 
 setup(
     name='trailscraper',
-    version='0.1',
+    version='0.2',
     description='A command-line tool to get valuable information out of AWS CloudTrail',
     long_description=readme + '\n\n' + changelog,
     url='http://github.com/flosell/trailscraper',

tests/cloudtrail/cloudtrail_test.py

@@ -1,19 +1,50 @@
 from tests.test_utils_testdata import cloudtrail_data, cloudtrail_data_dir
-from trailscraper.cloudtrail import _parse_records_from_gzipped_file, load_from_dir
 from trailscraper.cloudtrail import Record
+from trailscraper.cloudtrail import _parse_records_from_gzipped_file, load_from_dir, _parse_record, \
+    _parse_records
 
 
 def test_parse_records_from_gzipped_file():
     parsed_records = _parse_records_from_gzipped_file(cloudtrail_data("someRecords.json.gz"))
     assert parsed_records == [
-        Record("autoscaling.amazonaws.com", "DescribeLaunchConfigurations"),
-        Record("sts.amazonaws.com", "AssumeRole", ["arn:aws:iam::111111111111:role/someRole"])
+        Record("autoscaling.amazonaws.com", "DescribeLaunchConfigurations",
+               assumed_role_arn="arn:aws:iam::111111111111:role/someRole"),
+        Record("sts.amazonaws.com", "AssumeRole",
+               resource_arns=["arn:aws:iam::111111111111:role/someRole"])
     ]
 
 
 def test_load_all_gzipped_files_from_dir():
     records = load_from_dir(cloudtrail_data_dir())
     assert records == [
-        Record("autoscaling.amazonaws.com", "DescribeLaunchConfigurations"),
-        Record("sts.amazonaws.com", "AssumeRole", ["arn:aws:iam::111111111111:role/someRole"])
+        Record("autoscaling.amazonaws.com", "DescribeLaunchConfigurations",
+               assumed_role_arn="arn:aws:iam::111111111111:role/someRole"),
+        Record("sts.amazonaws.com", "AssumeRole",
+               resource_arns=["arn:aws:iam::111111111111:role/someRole"])
     ]
+
+
+def test_parse_record_should_be_able_to_cope_with_missing_type():
+    assert _parse_record({'userIdentity': {'accountId': '111111111111'},
+                          'eventSource': 'kms.amazonaws.com',
+                          'eventName': 'DeleteKey'}) == \
+           Record('kms.amazonaws.com', 'DeleteKey')
+
+
+def test_parse_record_should_be_able_to_cope_with_missing_session_context_in_assumed_role():
+    assert _parse_record({'eventVersion': '1.05',
+                          'userIdentity': {'type': 'AssumedRole', 'principalId': 'some-key:some-user',
+                                           'arn': 'arn:aws:sts::111111111111:assumed-role/some-role/some-user',
+                                           'accountId': '111111111111'},
+                          'eventSource': 'signin.amazonaws.com',
+                          'eventName': 'RenewRole'}) == \
+           Record('signin.amazonaws.com', 'RenewRole')
+
+
+def test_parse_records_should_ignore_records_that_cant_be_parsed():
+    assert _parse_records([{},
+                           {'eventVersion': '1.05',
+                            'userIdentity': {'type': 'SomeType'},
+                            'eventSource': 'someSource',
+                            'eventName': 'SomeEvent'}]) == \
+           [Record('someSource', 'SomeEvent')]

tests/cloudtrail/record_test.py

@@ -2,20 +2,34 @@
 
 
 def test_should_have_a_string_representation():
-    assert str(Record("sts.amazonaws.com", "AssumeRole",["arn:aws:iam::111111111111:role/someRole"])) == \
+    assert str(Record("sts.amazonaws.com", "AssumeRole",
+                      resource_arns = ["arn:aws:iam::111111111111:role/someRole"])) == \
            "Record(event_source=sts.amazonaws.com event_name=AssumeRole resource_arns=['arn:aws:iam::111111111111:role/someRole'])"
 
 
 def test_should_know_about_equality():
     assert Record("sts.amazonaws.com", "AssumeRole") == Record("sts.amazonaws.com", "AssumeRole")
     assert Record("sts.amazonaws.com", "AssumeRole",[]) == Record("sts.amazonaws.com", "AssumeRole")
-    assert Record("sts.amazonaws.com", "AssumeRole",["arn:aws:iam::111111111111:role/someRole"]) == \
-           Record("sts.amazonaws.com", "AssumeRole",["arn:aws:iam::111111111111:role/someRole"])
+    assert Record("sts.amazonaws.com", "AssumeRole",
+                  resource_arns = ["arn:aws:iam::111111111111:role/someRole"]) == \
+           Record("sts.amazonaws.com", "AssumeRole",
+                  resource_arns = ["arn:aws:iam::111111111111:role/someRole"])
+    assert Record("sts.amazonaws.com", "AssumeRole",
+                  assumed_role_arn = "arn:aws:iam::111111111111:role/someRole") == \
+           Record("sts.amazonaws.com", "AssumeRole",
+                  assumed_role_arn = "arn:aws:iam::111111111111:role/someRole")
+
     assert Record("sts.amazonaws.com", "AssumeRole") != Record("sts.amazonaws.com", "AssumeRoles")
     assert Record("sts.amazonaws.com", "AssumeRole") != Record("ec2.amazonaws.com", "AssumeRole")
     assert Record("sts.amazonaws.com", "AssumeRole") != Record("ec2.amazonaws.com", "DescribeInstances")
-    assert Record("sts.amazonaws.com", "AssumeRole", ["arn:aws:iam::111111111111:role/someRole"]) != \
+    assert Record("sts.amazonaws.com", "AssumeRole",
+                  resource_arns = ["arn:aws:iam::111111111111:role/someRole"]) != \
            Record("sts.amazonaws.com", "AssumeRole", ["arn:aws:iam::222222222222:role/someRole"])
+    assert Record("sts.amazonaws.com", "AssumeRole",
+                  assumed_role_arn = "arn:aws:iam::111111111111:role/someRole") != \
+           Record("sts.amazonaws.com", "AssumeRole",
+                  assumed_role_arn = "arn:aws:iam::111111111111:role/someOtherRole")
+
 
 
 def test_should_be_hashable():

tests/integration/cli_test.py

@@ -40,3 +40,27 @@ def test_should_output_an_iam_policy_for_a_set_of_cloudtrail_records():
     "Version": "2012-10-17"
 }
 '''
+
+
+def test_should_filter_for_assumed_role_arn():
+    runner = CliRunner()
+    result = runner.invoke(cli.root_group, args=["generate-policy", "--log-dir", cloudtrail_data_dir(),
+                                                 "--filter-assumed-role-arn",
+                                                 "arn:aws:iam::111111111111:role/someRole"])
+    assert result.output == '''\
+{
+    "Statement": [
+        {
+            "Action": [
+                "autoscaling:DescribeLaunchConfigurations"
+            ],
+            "Effect": "Allow",
+            "Resource": [
+                "*"
+            ]
+        }
+    ],
+    "Version": "2012-10-17"
+}
+'''
+    assert result.exit_code == 0

trailscraper/__init__.py

@@ -2,4 +2,4 @@
 
 __author__ = """Florian Sellmayr"""
 __email__ = 'florian.sellmayr@gmail.com'
-__version__ = '0.1.0'
+__version__ = '0.2.0'

trailscraper/cli.py

@@ -20,7 +20,6 @@ def root_group(verbose):
         logging.getLogger('s3transfer').setLevel(logging.INFO)
 
 
-
 @click.command()
 @click.option('--past-days', default=0, help='How many days to look into the past. 0 means today')
 @click.option('--bucket', required=True, help='The S3 bucket that contains cloud-trail logs')
@@ -38,12 +37,14 @@ def download(past_days, bucket, prefix, account_id, region, log_dir):
 
 @click.command("generate-policy")
 @click.option('--log-dir', default="~/.trailscraper/logs", type=click.Path(), help='Where to put logfiles')
-def generate_policy(log_dir):
+@click.option('--filter-assumed-role-arn', multiple=True,
+              help='only consider events from this role (can be used multiple times)')
+def generate_policy(log_dir, filter_assumed_role_arn):
     """Generates a policy that allows the events covered in the log-dir"""
     log_dir = os.path.expanduser(log_dir)
     records = load_from_dir(log_dir)
 
-    policy = generate_policy_from_records(records)
+    policy = generate_policy_from_records(records, filter_assumed_role_arn)
 
     click.echo(render_policy(policy))
 

trailscraper/cloudtrail.py

@@ -8,10 +8,11 @@
 class Record():
     """Represents a CloudTrail record"""
 
-    def __init__(self, event_source, event_name, resource_arns=None):
+    def __init__(self, event_source, event_name, resource_arns=None, assumed_role_arn=None):
         self.event_source = event_source
         self.event_name = event_name
         self.resource_arns = resource_arns or ["*"]
+        self.assumed_role_arn = assumed_role_arn
 
     def __repr__(self):
         return f"Record(event_source={self.event_source} event_name={self.event_name} " \
@@ -21,12 +22,16 @@ def __eq__(self, other):
         if isinstance(other, self.__class__):
             return self.event_source == other.event_source and \
                    self.event_name == other.event_name and \
-                   self.resource_arns == other.resource_arns
+                   self.resource_arns == other.resource_arns and \
+                   self.assumed_role_arn == other.assumed_role_arn
 
         return False
 
     def __hash__(self):
-        return hash((self.event_source, self.event_name, tuple(self.resource_arns)))
+        return hash((self.event_source,
+                     self.event_name,
+                     tuple(self.resource_arns),
+                     self.assumed_role_arn))
 
     def __ne__(self, other):
         return not self.__eq__(other)
@@ -38,10 +43,29 @@ def _resource_arns(json_record):
     return arns
 
 
-def _mk_record(json_record):
-    return Record(json_record['eventSource'],
-                  json_record['eventName'],
-                  _resource_arns(json_record))
+def _assumed_role_arn(json_record):
+    user_identity = json_record['userIdentity']
+    if 'type' in user_identity \
+            and user_identity['type'] == 'AssumedRole' \
+            and 'sessionContext' in user_identity:
+        return user_identity['sessionContext']['sessionIssuer']['arn']
+    return None
+
+
+def _parse_record(json_record):
+    try:
+        return Record(json_record['eventSource'],
+                      json_record['eventName'],
+                      resource_arns=_resource_arns(json_record),
+                      assumed_role_arn=_assumed_role_arn(json_record))
+    except KeyError as error:
+        logging.warning("Could not parse %s: %s", json_record, error)
+        return None
+
+
+def _parse_records(json_records):
+    parsed_records = [_parse_record(record) for record in json_records]
+    return [r for r in parsed_records if r is not None]
 
 
 def _parse_records_from_gzipped_file(filename):
@@ -50,7 +74,8 @@ def _parse_records_from_gzipped_file(filename):
 
     with gzip.open(filename, 'rb') as file:
         json_data = json.load(file)
-        return [_mk_record(record) for record in json_data['Records']]
+        records = json_data['Records']
+        return _parse_records(records)
 
 
 def load_from_dir(log_dir):

trailscraper/policy_generator.py

@@ -41,15 +41,18 @@ def _combine_statements_with_the_same_actions(statements):
     return combined_statements
 
 
-def generate_policy_from_records(records):
+def generate_policy_from_records(records, arns_to_filter_for=None):
     """Generates a policy from a set of records"""
 
+    if arns_to_filter_for is None:
+        arns_to_filter_for = []
+
     statements = [
         Statement(
             Effect="Allow",
             Action=[Action(_source_to_iam_prefix(record.event_source), record.event_name)],
             Resource=sorted(record.resource_arns)
-        ) for record in records
+        ) for record in records if (record.assumed_role_arn in arns_to_filter_for) or (len(arns_to_filter_for) == 0)
     ]
 
     combined_statements = _combine_statements_with_the_same_actions(