CoCalc -- test_validation.py

GitHub Repository: aws/aws-cli
Path: blob/develop/tests/unit/customizations/cloudtrail/test_validation.py
²⁶³⁰ views
1
# Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2
#
3
# Licensed under the Apache License, Version 2.0 (the 'License'). You
4
# may not use this file except in compliance with the License. A copy of
5
# the License is located at
6
#
7
#     http://aws.amazon.com/apache2.0/
8
#
9
# or in the 'license' file accompanying this file. This file is
10
# distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11
# ANY KIND, either express or implied. See the License for the specific
12
# language governing permissions and limitations under the License.
13
import binascii
14
import base64
15
import hashlib
16
import json
17
import gzip
18
from datetime import datetime, timedelta
19
from dateutil import parser, tz
20

21
import rsa
22
from argparse import Namespace
23

24
from awscli.testutils import BaseAWSCommandParamsTest
25
from awscli.customizations.cloudtrail.validation import DigestError, \
26
    extract_digest_key_date, normalize_date, format_date, DigestProvider, \
27
    DigestTraverser, create_digest_traverser, PublicKeyProvider, \
28
    Sha256RSADigestValidator, DATE_FORMAT, CloudTrailValidateLogs, \
29
    parse_date, assert_cloudtrail_arn_is_valid, DigestSignatureError, \
30
    InvalidDigestFormat, S3ClientProvider
31
from awscli.compat import BytesIO
32
from botocore.exceptions import ClientError
33
from awscli.testutils import mock, unittest
34
from awscli.schema import ParameterRequiredError
35

36

37
START_DATE = parser.parse('20140810T000000Z')
38
END_DATE = parser.parse('20150810T000000Z')
39
TEST_ACCOUNT_ID = '123456789012'
40
TEST_TRAIL_ARN = 'arn:aws:cloudtrail:us-east-1:%s:trail/foo' % TEST_ACCOUNT_ID
41
VALID_TEST_KEY = ('MIIBCgKCAQEAn11L2YZ9h7onug2ILi1MWyHiMRsTQjfWE+pHVRLk1QjfW'
42
                  'hirG+lpOa8NrwQ/r7Ah5bNL6HepznOU9XTDSfmmnP97mqyc7z/upfZdS/'
43
                  'AHhYcGaz7n6Wc/RRBU6VmiPCrAUojuSk6/GjvA8iOPFsYDuBtviXarvuL'
44
                  'PlrT9kAd4Lb+rFfR5peEgBEkhlzc5HuWO7S0y+KunqxX6jQBnXGMtxmPB'
45
                  'PP0FylgWGNdFtks/4YSKcgqwH0YDcawP9GGGDAeCIqPWIXDLG1jOjRRzW'
46
                  'fCmD0iJUkz8vTsn4hq/5ZxRFE7UBAUiVcGbdnDdvVfhF9C3dQiDq3k7ad'
47
                  'QIziLT0cShgQIDAQAB')
48
TEST_ORGANIZATION_ACCOUNT_ID = '987654321098'
49
TEST_ORGANIZATION_ID = 'o-12345'
50

51

52
def create_mock_key_provider(key_list):
53
    """Creates a mock key provider that yields keys for each in key_list"""
54
    public_keys = {}
55
    for k in key_list:
56
        public_keys[k] = {'Fingerprint': k,
57
                          'Value': 'ffaa00'}
58
    key_provider = mock.Mock()
59
    key_provider.get_public_keys.return_value = public_keys
60
    return key_provider
61

62

63
def create_scenario(actions, logs=None):
64
    """Creates a scenario for a stack of actions
65

66
    Each action can be "gap" meaning there is no previous link, "invalid"
67
    meaning we should simulate an invalid digest, "missing" meaning we
68
    should simulate a digest is missing from S3, "bucket_change" meaning
69
    it is a link but the bucket is different than the previous bucket.
70
    Values are popped one by one off of the list until a terminal "gap"
71
    action is found.
72
    """
73
    keys = [str(i) for i in range(len(actions))]
74
    key_provider = create_mock_key_provider(keys)
75
    digest_provider = MockDigestProvider(actions, logs)
76
    digest_validator = mock.Mock()
77

78
    def validate(bucket, key, public_key, digest_data, digest_str):
79
        if '_invalid' in digest_data:
80
            raise DigestError('invalid error')
81

82
    digest_validator.validate = validate
83
    return key_provider, digest_provider, digest_validator
84

85

86
def collecting_callback():
87
    """Create and return a callback and a list populated with call args"""
88
    calls = []
89

90
    def cb(**kwargs):
91
        calls.append(kwargs)
92

93
    return cb, calls
94

95

96
class MockDigestProvider(object):
97
    def __init__(self, actions, logs=None):
98
        self.logs = logs or []
99
        self.actions = actions
100
        self.calls = {'fetch_digest': [], 'load_digest_keys_in_range': []}
101
        self.digests = []
102
        for i in range(len(self.actions)):
103
            self.digests.append(self.get_key_at_position(i))
104

105
    def get_key_at_position(self, position):
106
        dt = START_DATE + timedelta(hours=position)
107
        key = ('AWSLogs/{account}/CloudTrail-Digest/us-east-1/{ymd}/{account}_'
108
               'CloudTrail-Digest_us-east-1_foo_us-east-1_{date}.json.gz')
109
        return key.format(
110
            account=TEST_ACCOUNT_ID,
111
            ymd=dt.strftime('%Y/%m/%d'),
112
            date=dt.strftime(DATE_FORMAT))
113

114
    @staticmethod
115
    def create_digest(fingerprint, start_date, key, bucket, next_key=None,
116
                      next_bucket=None, logs=None):
117
        digest_end_date = start_date + timedelta(hours=1, minutes=30)
118
        return {'digestPublicKeyFingerprint': fingerprint,
119
                'digestEndTime': digest_end_date.strftime(DATE_FORMAT),
120
                'digestStartTime': start_date.strftime(DATE_FORMAT),
121
                'previousDigestS3Bucket': next_bucket,
122
                'previousDigestS3Object': next_key,
123
                'digestS3Bucket': bucket,
124
                'digestS3Object': key,
125
                'awsAccountId': TEST_ACCOUNT_ID,
126
                'previousDigestSignature': 'abcd',
127
                'logFiles': logs or []}
128

129
    @staticmethod
130
    def create_link(key, next_key, next_bucket, position, action, logs,
131
                    bucket):
132
        """Creates a link in a digest chain for testing."""
133
        digest_logs = []
134
        if len(logs) > position:
135
            digest_logs = logs[position]
136
        end_date = parse_date(extract_digest_key_date(key))
137
        # gap actions have no previous link.
138
        if action == 'gap':
139
            digest = MockDigestProvider.create_digest(
140
                key=key, bucket=bucket, fingerprint=str(position),
141
                start_date=end_date, logs=digest_logs)
142
        else:
143
            digest = MockDigestProvider.create_digest(
144
                key=key, bucket=bucket, fingerprint=str(position),
145
                start_date=end_date, next_bucket=next_bucket, next_key=next_key,
146
                logs=digest_logs)
147
            # Mark the digest as invalid if specified in the action.
148
            if action == 'invalid':
149
                digest['_invalid'] = True
150
        return digest, json.dumps(digest)
151

152
    def load_digest_keys_in_range(self, bucket, prefix, start_date, end_date):
153
        self.calls['load_digest_keys_in_range'].append(locals())
154
        return list(self.digests)
155

156
    def fetch_digest(self, bucket, key):
157
        self.calls['fetch_digest'].append(key)
158
        position = self.digests.index(key)
159
        action = self.actions[position]
160
        # Simulate a digest missing from S3
161
        if action == 'missing':
162
            raise ClientError(
163
                {'Error': {'Code': 'NoSuchKey', 'Message': 'foo'}},
164
                'GetObject')
165
        next_key = self.get_key_at_position(position - 1)
166
        next_bucket = int(bucket)
167
        if action == 'bucket_change':
168
            next_bucket += 1
169
        return self.create_link(key, next_key, str(next_bucket), position,
170
                                action, self.logs, bucket)
171

172

173
class TestValidation(unittest.TestCase):
174
    def test_formats_dates(self):
175
        date = datetime(2015, 8, 21, tzinfo=tz.tzutc())
176
        self.assertEqual('20150821T000000Z', format_date(date))
177

178
    def test_parses_dates_with_better_error_message(self):
179
        try:
180
            parse_date('foo')
181
            self.fail('Should have failed to parse')
182
        except ValueError as e:
183
            self.assertIn('Unable to parse date value: foo', str(e))
184

185
    def test_parses_dates(self):
186
        date = parse_date('August 25, 2015 00:00:00 UTC')
187
        self.assertEqual(date, datetime(2015, 8, 25, tzinfo=tz.tzutc()))
188

189
    def test_ensures_cloudtrail_arns_are_valid(self):
190
        try:
191
            assert_cloudtrail_arn_is_valid('foo:bar:baz')
192
            self.fail('Should have failed')
193
        except ValueError as e:
194
            self.assertIn('Invalid trail ARN provided: foo:bar:baz', str(e))
195

196
    def test_ensures_cloudtrail_arns_are_valid_when_missing_resource(self):
197
        try:
198
            assert_cloudtrail_arn_is_valid(
199
                'arn:aws:cloudtrail:us-east-1:%s:foo' % TEST_ACCOUNT_ID)
200
            self.fail('Should have failed')
201
        except ValueError as e:
202
            self.assertIn('Invalid trail ARN provided', str(e))
203

204
    def test_allows_valid_arns(self):
205
        assert_cloudtrail_arn_is_valid(
206
            'arn:aws:cloudtrail:us-east-1:%s:trail/foo' % TEST_ACCOUNT_ID)
207

208
    def test_normalizes_date_timezones(self):
209
        date = datetime(2015, 8, 21, tzinfo=tz.tzlocal())
210
        normalized = normalize_date(date)
211
        self.assertEqual(tz.tzutc(), normalized.tzinfo)
212

213
    def test_extracts_dates_from_digest_keys(self):
214
        arn = ('AWSLogs/{account}/CloudTrail-Digest/us-east-1/2015/08/'
215
               '16/{account}_CloudTrail-Digest_us-east-1_foo_us-east-1_'
216
               '20150816T230550Z.json.gz').format(account=TEST_ACCOUNT_ID)
217
        self.assertEqual('20150816T230550Z', extract_digest_key_date(arn))
218

219
    def test_creates_traverser(self):
220
        mock_s3_provider = mock.Mock()
221
        traverser = create_digest_traverser(
222
            trail_arn=TEST_TRAIL_ARN, cloudtrail_client=mock.Mock(),
223
            organization_client=mock.Mock(),
224
            trail_source_region='us-east-1',
225
            s3_client_provider=mock_s3_provider,
226
            bucket='bucket', prefix='prefix')
227
        self.assertEqual('bucket', traverser.starting_bucket)
228
        self.assertEqual('prefix', traverser.starting_prefix)
229
        digest_provider = traverser.digest_provider
230
        self.assertEqual('us-east-1', digest_provider.trail_home_region)
231
        self.assertEqual('foo', digest_provider.trail_name)
232

233
    def test_creates_traverser_account_id(self):
234
        mock_s3_provider = mock.Mock()
235
        traverser = create_digest_traverser(
236
            trail_arn=TEST_TRAIL_ARN, cloudtrail_client=mock.Mock(),
237
            organization_client=mock.Mock(),
238
            trail_source_region='us-east-1',
239
            s3_client_provider=mock_s3_provider,
240
            bucket='bucket', prefix='prefix',
241
            account_id=TEST_ORGANIZATION_ACCOUNT_ID)
242
        self.assertEqual('bucket', traverser.starting_bucket)
243
        self.assertEqual('prefix', traverser.starting_prefix)
244
        digest_provider = traverser.digest_provider
245
        self.assertEqual('us-east-1', digest_provider.trail_home_region)
246
        self.assertEqual('foo', digest_provider.trail_name)
247
        self.assertEqual(
248
            TEST_ORGANIZATION_ACCOUNT_ID, digest_provider.account_id)
249

250
    def test_creates_traverser_and_gets_trail_by_arn(self):
251
        cloudtrail_client = mock.Mock()
252
        cloudtrail_client.describe_trails.return_value = {'trailList': [
253
            {'TrailARN': TEST_TRAIL_ARN,
254
             'S3BucketName': 'bucket', 'S3KeyPrefix': 'prefix',
255
             'IsOrganizationTrail': False}
256
        ]}
257
        traverser = create_digest_traverser(
258
            trail_arn=TEST_TRAIL_ARN, trail_source_region='us-east-1',
259
            cloudtrail_client=cloudtrail_client,
260
            organization_client=mock.Mock(),
261
            s3_client_provider=mock.Mock())
262
        self.assertEqual('bucket', traverser.starting_bucket)
263
        self.assertEqual('prefix', traverser.starting_prefix)
264
        digest_provider = traverser.digest_provider
265
        self.assertEqual('us-east-1', digest_provider.trail_home_region)
266
        self.assertEqual('foo', digest_provider.trail_name)
267
        self.assertEqual(TEST_ACCOUNT_ID, digest_provider.account_id)
268

269
    def test_create_traverser_organizational_trail_not_launched(self):
270
        cloudtrail_client = mock.Mock()
271
        cloudtrail_client.describe_trails.return_value = {'trailList': [
272
            {'TrailARN': TEST_TRAIL_ARN,
273
             'S3BucketName': 'bucket', 'S3KeyPrefix': 'prefix'}
274
        ]}
275
        traverser = create_digest_traverser(
276
            trail_arn=TEST_TRAIL_ARN, trail_source_region='us-east-1',
277
            cloudtrail_client=cloudtrail_client,
278
            organization_client=mock.Mock(),
279
            s3_client_provider=mock.Mock())
280
        self.assertEqual('bucket', traverser.starting_bucket)
281
        self.assertEqual('prefix', traverser.starting_prefix)
282
        digest_provider = traverser.digest_provider
283
        self.assertEqual('us-east-1', digest_provider.trail_home_region)
284
        self.assertEqual('foo', digest_provider.trail_name)
285
        self.assertEqual(TEST_ACCOUNT_ID, digest_provider.account_id)
286

287
    def test_creates_traverser_and_gets_trail_by_arn_s3_bucket_specified(self):
288
        cloudtrail_client = mock.Mock()
289
        traverser = create_digest_traverser(
290
            trail_arn=TEST_TRAIL_ARN, trail_source_region='us-east-1',
291
            cloudtrail_client=cloudtrail_client,
292
            organization_client=mock.Mock(),
293
            s3_client_provider=mock.Mock(),
294
            bucket="bucket")
295
        self.assertEqual('bucket', traverser.starting_bucket)
296
        digest_provider = traverser.digest_provider
297
        self.assertEqual('us-east-1', digest_provider.trail_home_region)
298
        self.assertEqual('foo', digest_provider.trail_name)
299
        self.assertEqual(TEST_ACCOUNT_ID, digest_provider.account_id)
300

301
    def test_creates_traverser_and_gets_organization_id(self):
302
        cloudtrail_client = mock.Mock()
303
        cloudtrail_client.describe_trails.return_value = {'trailList': [
304
            {'TrailARN': TEST_TRAIL_ARN,
305
             'S3BucketName': 'bucket', 'S3KeyPrefix': 'prefix',
306
             'IsOrganizationTrail': True}
307
        ]}
308
        organization_client = mock.Mock()
309
        organization_client.describe_organization.return_value = {
310
            "Organization": {
311
                "MasterAccountId": TEST_ACCOUNT_ID,
312
                "Id": TEST_ORGANIZATION_ID,
313
            }
314
        }
315
        traverser = create_digest_traverser(
316
            trail_arn=TEST_TRAIL_ARN, trail_source_region='us-east-1',
317
            cloudtrail_client=cloudtrail_client,
318
            organization_client=organization_client,
319
            s3_client_provider=mock.Mock(), account_id=TEST_ACCOUNT_ID)
320
        self.assertEqual('bucket', traverser.starting_bucket)
321
        self.assertEqual('prefix', traverser.starting_prefix)
322
        digest_provider = traverser.digest_provider
323
        self.assertEqual('us-east-1', digest_provider.trail_home_region)
324
        self.assertEqual('foo', digest_provider.trail_name)
325
        self.assertEqual(TEST_ORGANIZATION_ID, digest_provider.organization_id)
326

327
    def test_creates_traverser_organization_trail_missing_account_id(self):
328
        cloudtrail_client = mock.Mock()
329
        cloudtrail_client.describe_trails.return_value = {'trailList': [
330
            {'TrailARN': TEST_TRAIL_ARN,
331
             'S3BucketName': 'bucket', 'S3KeyPrefix': 'prefix',
332
             'IsOrganizationTrail': True}
333
        ]}
334
        organization_client = mock.Mock()
335
        organization_client.describe_organization.return_value = {
336
            "Organization": {
337
                "MasterAccountId": TEST_ACCOUNT_ID,
338
                "Id": TEST_ORGANIZATION_ID,
339
            }
340
        }
341
        with self.assertRaises(ParameterRequiredError):
342
            create_digest_traverser(
343
                trail_arn=TEST_TRAIL_ARN, trail_source_region='us-east-1',
344
                cloudtrail_client=cloudtrail_client,
345
                organization_client=organization_client,
346
                s3_client_provider=mock.Mock())
347

348

349
class TestPublicKeyProvider(unittest.TestCase):
350
    def test_returns_public_key_in_range(self):
351
        cloudtrail_client = mock.Mock()
352
        cloudtrail_client.list_public_keys.return_value = {'PublicKeyList': [
353
            {'Fingerprint': 'a', 'OtherData': 'a', 'Value': 'a'},
354
            {'Fingerprint': 'b', 'OtherData': 'b', 'Value': 'b'},
355
            {'Fingerprint': 'c', 'OtherData': 'c', 'Value': 'c'},
356
        ]}
357
        provider = PublicKeyProvider(cloudtrail_client)
358
        start_date = START_DATE
359
        end_date = start_date + timedelta(days=2)
360
        keys = provider.get_public_keys(start_date, end_date)
361
        self.assertEqual({
362
            'a': {'Fingerprint': 'a', 'OtherData': 'a', 'Value': 'a'},
363
            'b': {'Fingerprint': 'b', 'OtherData': 'b', 'Value': 'b'},
364
            'c': {'Fingerprint': 'c', 'OtherData': 'c', 'Value': 'c'},
365
        }, keys)
366
        cloudtrail_client.list_public_keys.assert_has_calls(
367
            [mock.call(EndTime=end_date, StartTime=start_date)])
368

369

370
class TestSha256RSADigestValidator(unittest.TestCase):
371
    def setUp(self):
372
        self._digest_data = {'digestStartTime': 'baz',
373
                             'digestEndTime': 'foo',
374
                             'awsAccountId': 'account',
375
                             'digestPublicKeyFingerprint': 'abc',
376
                             'digestS3Bucket': 'bucket',
377
                             'digestS3Object': 'object',
378
                             'previousDigestSignature': 'xyz'}
379
        self._inflated_digest = json.dumps(self._digest_data).encode()
380
        self._digest_data['_signature'] = 'aeff'
381

382
    def test_validates_digests(self):
383
        (public_key, private_key) = rsa.newkeys(512)
384
        sha256_hash = hashlib.sha256(self._inflated_digest)
385
        string_to_sign = "%s\n%s/%s\n%s\n%s" % (
386
            self._digest_data['digestEndTime'],
387
            self._digest_data['digestS3Bucket'],
388
            self._digest_data['digestS3Object'],
389
            sha256_hash.hexdigest(),
390
            self._digest_data['previousDigestSignature'])
391
        signature = rsa.sign(string_to_sign.encode(), private_key, 'SHA-256')
392
        self._digest_data['_signature'] = binascii.hexlify(signature)
393
        validator = Sha256RSADigestValidator()
394
        public_key_b64 = base64.b64encode(public_key.save_pkcs1(format='DER'))
395
        validator.validate('b', 'k', public_key_b64, self._digest_data,
396
                           self._inflated_digest)
397

398
    def test_does_not_expose_underlying_key_decoding_error(self):
399
        validator = Sha256RSADigestValidator()
400
        try:
401
            validator.validate(
402
                'b', 'k', 'YQo=', self._digest_data, 'invalid'.encode())
403
            self.fail('Should have failed')
404
        except DigestError as e:
405
            self.assertEqual(('Digest file\ts3://b/k\tINVALID: Unable to load '
406
                              'PKCS #1 key with fingerprint abc'), str(e))
407

408
    def test_does_not_expose_underlying_validation_error(self):
409
        validator = Sha256RSADigestValidator()
410
        try:
411
            validator.validate(
412
                'b', 'k', VALID_TEST_KEY, self._digest_data,
413
                'invalid'.encode())
414
            self.fail('Should have failed')
415
        except DigestSignatureError as e:
416
            self.assertEqual(('Digest file\ts3://b/k\tINVALID: signature '
417
                              'verification failed'), str(e))
418

419
    def test_properly_signs_when_no_previous_signature(self):
420
        validator = Sha256RSADigestValidator()
421
        digest_data = {
422
            'digestEndTime': 'a',
423
            'digestS3Bucket': 'b',
424
            'digestS3Object': 'c',
425
            'previousDigestSignature': None}
426
        signed = validator._create_string_to_sign(digest_data, 'abc'.encode())
427
        self.assertEqual(
428
            ('a\nb/c\nba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff6'
429
             '1f20015ad\nnull').encode(), signed)
430

431

432
class TestDigestProvider(BaseAWSCommandParamsTest):
433
    def _fake_key(self, date):
434
        parsed = parser.parse(date)
435
        return ('prefix/AWSLogs/{account}/CloudTrail-Digest/us-east-1/{year}/'
436
                '{month}/{day}/{account}_CloudTrail-Digest_us-east-1_foo_'
437
                'us-east-1_{date}.json.gz').format(date=date, year=parsed.year,
438
                                                   month=parsed.month,
439
                                                   account=TEST_ACCOUNT_ID,
440
                                                   day=parsed.day)
441

442
    def _get_mock_provider(self, s3_client):
443
        mock_s3_client_provider = mock.Mock()
444
        mock_s3_client_provider.get_client.return_value = s3_client
445
        return DigestProvider(
446
            mock_s3_client_provider, TEST_ACCOUNT_ID, 'foo', 'us-east-1')
447

448
    def test_initializes_public_properties(self):
449
        client = mock.Mock()
450
        provider = DigestProvider(client, TEST_ACCOUNT_ID, 'foo', 'us-east-1')
451
        self.assertEqual(TEST_ACCOUNT_ID, provider.account_id)
452
        self.assertEqual('foo', provider.trail_name)
453
        self.assertEqual('us-east-1', provider.trail_home_region)
454

455
    def test_returns_digests_in_range(self):
456
        s3_client = self.driver.session.create_client('s3')
457
        keys = [self._fake_key(format_date(START_DATE - timedelta(days=1))),
458
                self._fake_key(format_date(START_DATE + timedelta(days=1))),
459
                self._fake_key(format_date(START_DATE + timedelta(days=2))),
460
                self._fake_key(format_date(START_DATE + timedelta(days=3))),
461
                self._fake_key(format_date(END_DATE + timedelta(hours=1))),
462
                self._fake_key(format_date(END_DATE + timedelta(days=1)))]
463
        # Create a key that looks similar but for a different trail.
464
        bad_name = keys[3].replace('foo', 'baz')
465
        # Create a key that looks similar but is from a different trail source
466
        # region (e.g., CloudTrail-Digest/us-west-2).
467
        bad_region = keys[3].replace(
468
            'CloudTrail-Digest/us-east-1', 'CloudTrail-Digest/us-west-2')
469
        bad_region = bad_region.replace(
470
            'CloudTrail-Digest_us-east-1', 'CloudTrail-Digest_us-west-2')
471
        self.parsed_responses = [
472
            {"Contents": [{"Key": keys[0]},        # skip (date <)
473
                          {"Key": keys[1]},
474
                          {"Key": keys[2]},
475
                          {"Key": 'foo/baz/bar'},  # skip (regex (bogus))
476
                          {"Key": bad_name},       # skip (regex (trail name))
477
                          {"Key": bad_region},     # skip (regex (source))
478
                          {"Key": keys[3]},
479
                          {"Key": keys[4]},        # hour is +1, but keep
480
                          {"Key": keys[5]}]}]      # skip (date >)
481
        self.patch_make_request()
482
        provider = self._get_mock_provider(s3_client)
483
        digests = provider.load_digest_keys_in_range(
484
            'foo', 'prefix', START_DATE, END_DATE)
485
        self.assertNotIn(bad_name, digests)
486
        self.assertNotIn(bad_region, digests)
487
        self.assertEqual(keys[1], digests[0])
488
        self.assertEqual(keys[2], digests[1])
489
        self.assertEqual(keys[3], digests[2])
490
        self.assertEqual(keys[4], digests[3])
491

492
    def test_calls_list_objects_correctly(self):
493
        s3_client = mock.Mock()
494
        mock_paginate = s3_client.get_paginator.return_value.paginate
495
        mock_search = mock_paginate.return_value.search
496
        mock_search.return_value = []
497
        provider = self._get_mock_provider(s3_client)
498
        provider.load_digest_keys_in_range(
499
            '1', 'prefix', START_DATE, END_DATE)
500
        marker = ('prefix/AWSLogs/{account}/CloudTrail-Digest/us-east-1/'
501
                  '2014/08/09/{account}_CloudTrail-Digest_us-east-1_foo_'
502
                  'us-east-1_20140809T235900Z.json.gz')
503
        prefix = 'prefix/AWSLogs/{account}/CloudTrail-Digest/us-east-1'
504
        mock_paginate.assert_called_once_with(
505
            Bucket='1',
506
            Marker=marker.format(account=TEST_ACCOUNT_ID),
507
            Prefix=prefix.format(account=TEST_ACCOUNT_ID))
508

509
    def test_calls_list_objects_correctly_org_trails(self):
510
        s3_client = mock.Mock()
511
        mock_s3_client_provider = mock.Mock()
512
        mock_paginate = s3_client.get_paginator.return_value.paginate
513
        mock_search = mock_paginate.return_value.search
514
        mock_search.return_value = []
515
        mock_s3_client_provider.get_client.return_value = s3_client
516
        provider = DigestProvider(
517
            mock_s3_client_provider, TEST_ORGANIZATION_ACCOUNT_ID,
518
            'foo', 'us-east-1', 'us-east-1',
519
            TEST_ORGANIZATION_ID)
520
        provider.load_digest_keys_in_range(
521
            '1', 'prefix', START_DATE, END_DATE)
522
        marker = (
523
            'prefix/AWSLogs/{organization_id}/{member_account}/'
524
            'CloudTrail-Digest/us-east-1/'
525
            '2014/08/09/{member_account}_CloudTrail-Digest_us-east-1_foo_'
526
            'us-east-1_20140809T235900Z.json.gz'
527
        )
528
        prefix = (
529
            'prefix/AWSLogs/{organization_id}/{member_account}/'
530
            'CloudTrail-Digest/us-east-1'
531
        )
532
        mock_paginate.assert_called_once_with(
533
            Bucket='1',
534
            Marker=marker.format(
535
                member_account=TEST_ORGANIZATION_ACCOUNT_ID,
536
                organization_id=TEST_ORGANIZATION_ID
537
            ),
538
            Prefix=prefix.format(
539
                member_account=TEST_ORGANIZATION_ACCOUNT_ID,
540
                organization_id=TEST_ORGANIZATION_ID
541
            )
542
        )
543

544
    def test_create_digest_prefix_without_key_prefix(self):
545
        mock_s3_client_provider = mock.Mock()
546
        provider = DigestProvider(
547
            mock_s3_client_provider, TEST_ACCOUNT_ID, 'foo', 'us-east-1')
548
        prefix = provider._create_digest_prefix(START_DATE, None)
549
        expected = 'AWSLogs/{account}/CloudTrail-Digest/us-east-1'.format(
550
            account=TEST_ACCOUNT_ID)
551
        self.assertEqual(expected, prefix)
552

553
    def test_create_digest_prefix_with_key_prefix(self):
554
        mock_s3_client_provider = mock.Mock()
555
        provider = DigestProvider(
556
            mock_s3_client_provider, TEST_ACCOUNT_ID, 'foo', 'us-east-1')
557
        prefix = provider._create_digest_prefix(START_DATE, 'my-prefix')
558
        expected = 'my-prefix/AWSLogs/{account}/CloudTrail-Digest/us-east-1'.format(
559
            account=TEST_ACCOUNT_ID)
560
        self.assertEqual(expected, prefix)
561

562
    def test_create_digest_prefix_org_trail(self):
563
        mock_s3_client_provider = mock.Mock()
564
        provider = DigestProvider(
565
            mock_s3_client_provider, TEST_ORGANIZATION_ACCOUNT_ID,
566
            'foo', 'us-east-1', 'us-east-1', TEST_ORGANIZATION_ID)
567
        prefix = provider._create_digest_prefix(START_DATE, None)
568
        expected = 'AWSLogs/{org}/{account}/CloudTrail-Digest/us-east-1'.format(
569
            org=TEST_ORGANIZATION_ID,
570
            account=TEST_ORGANIZATION_ACCOUNT_ID)
571
        self.assertEqual(expected, prefix)
572

573
    def test_create_digest_prefix_org_trail_with_key_prefix(self):
574
        mock_s3_client_provider = mock.Mock()
575
        provider = DigestProvider(
576
            mock_s3_client_provider, TEST_ORGANIZATION_ACCOUNT_ID,
577
            'foo', 'us-east-1', 'us-east-1', TEST_ORGANIZATION_ID)
578
        prefix = provider._create_digest_prefix(START_DATE, 'custom-prefix')
579
        expected = 'custom-prefix/AWSLogs/{org}/{account}/CloudTrail-Digest/us-east-1'.format(
580
            org=TEST_ORGANIZATION_ID,
581
            account=TEST_ORGANIZATION_ACCOUNT_ID)
582
        self.assertEqual(expected, prefix)
583

584
    def test_ensures_digest_has_proper_metadata(self):
585
        out = BytesIO()
586
        f = gzip.GzipFile(fileobj=out, mode="wb")
587
        f.write('{"foo":"bar"}'.encode())
588
        f.close()
589
        gzipped_data = out.getvalue()
590
        s3_client = mock.Mock()
591
        s3_client.get_object.return_value = {
592
            'Body': BytesIO(gzipped_data),
593
            'Metadata': {}}
594
        provider = self._get_mock_provider(s3_client)
595
        with self.assertRaises(DigestSignatureError):
596
            provider.fetch_digest('bucket', 'key')
597

598
    def test_ensures_digest_can_be_gzip_inflated(self):
599
        s3_client = mock.Mock()
600
        s3_client.get_object.return_value = {
601
            'Body': BytesIO('foo'.encode()),
602
            'Metadata': {}}
603
        provider = self._get_mock_provider(s3_client)
604
        with self.assertRaises(InvalidDigestFormat):
605
            provider.fetch_digest('bucket', 'key')
606

607
    def test_ensures_digests_can_be_json_parsed(self):
608
        json_str = '{{{'
609
        out = BytesIO()
610
        f = gzip.GzipFile(fileobj=out, mode="wb")
611
        f.write(json_str.encode())
612
        f.close()
613
        gzipped_data = out.getvalue()
614
        s3_client = mock.Mock()
615
        s3_client.get_object.return_value = {
616
            'Body': BytesIO(gzipped_data),
617
            'Metadata': {'signature': 'abc', 'signature-algorithm': 'SHA256'}}
618
        provider = self._get_mock_provider(s3_client)
619
        with self.assertRaises(InvalidDigestFormat):
620
            provider.fetch_digest('bucket', 'key')
621

622
    def test_fetches_digests(self):
623
        json_str = '{"foo":"bar"}'
624
        out = BytesIO()
625
        f = gzip.GzipFile(fileobj=out, mode="wb")
626
        f.write(json_str.encode())
627
        f.close()
628
        gzipped_data = out.getvalue()
629
        s3_client = mock.Mock()
630
        s3_client.get_object.return_value = {
631
            'Body': BytesIO(gzipped_data),
632
            'Metadata': {'signature': 'abc', 'signature-algorithm': 'SHA256'}}
633
        provider = self._get_mock_provider(s3_client)
634
        result = provider.fetch_digest('bucket', 'key')
635
        self.assertEqual({'foo': 'bar', '_signature': 'abc',
636
                          '_signature_algorithm': 'SHA256'}, result[0])
637
        self.assertEqual(json_str.encode(), result[1])
638

639

640
class TestDigestTraverser(unittest.TestCase):
641
    def test_initializes_with_default_validator(self):
642
        provider = mock.Mock()
643
        traverser = DigestTraverser(
644
            digest_provider=provider, starting_bucket='1',
645
            starting_prefix='baz', public_key_provider=mock.Mock())
646
        self.assertEqual('1', traverser.starting_bucket)
647
        self.assertEqual('baz', traverser.starting_prefix)
648
        self.assertEqual(provider, traverser.digest_provider)
649

650
    def test_ensures_public_keys_are_loaded(self):
651
        start_date = START_DATE
652
        end_date = END_DATE
653
        digest_provider = mock.Mock()
654
        key_provider = mock.Mock()
655
        key_provider.get_public_keys.return_value = []
656
        traverser = DigestTraverser(
657
            digest_provider=digest_provider, starting_bucket='1',
658
            starting_prefix='baz', public_key_provider=key_provider)
659
        digest_iter = traverser.traverse(start_date, end_date)
660
        with self.assertRaises(RuntimeError):
661
            next(digest_iter)
662
        key_provider.get_public_keys.assert_called_with(
663
            start_date, end_date)
664

665
    def test_ensures_public_key_is_found(self):
666
        start_date = START_DATE
667
        end_date = END_DATE
668
        key_name = end_date.strftime(DATE_FORMAT) + '.json.gz'
669
        region = 'us-west-2'
670
        digest_provider = mock.Mock()
671
        digest_provider.trail_home_region = region
672
        digest_provider.load_digest_keys_in_range.return_value = [key_name]
673
        digest_provider.fetch_digest.return_value = (
674
            {'digestEndTime': 'foo',
675
             'digestStartTime': 'foo',
676
             'awsAccountId': 'account',
677
             'digestPublicKeyFingerprint': 'abc',
678
             'digestS3Bucket': '1',
679
             'digestS3Object': key_name,
680
             'previousDigestSignature': 'xyz'},
681
            'abc'
682
        )
683
        key_provider = mock.Mock()
684
        key_provider.get_public_keys.return_value = [{'Fingerprint': 'a'}]
685
        on_invalid, calls = collecting_callback()
686
        traverser = DigestTraverser(
687
            digest_provider=digest_provider, starting_bucket='1',
688
            starting_prefix='baz', public_key_provider=key_provider,
689
            on_invalid=on_invalid)
690
        digest_iter = traverser.traverse(start_date, end_date)
691
        with self.assertRaises(StopIteration):
692
            next(digest_iter)
693
        self.assertEqual(1, len(calls))
694
        self.assertEqual(
695
            ('Digest file\ts3://1/%s\tINVALID: public key not '
696
             'found in region %s for fingerprint abc' % (key_name, region)),
697
            calls[0]['message'])
698

699
    def test_invokes_digest_validator(self):
700
        start_date = START_DATE
701
        end_date = END_DATE
702
        key_name = end_date.strftime(DATE_FORMAT) + '.json.gz'
703
        digest = {'digestPublicKeyFingerprint': 'a',
704
                  'digestS3Bucket': '1',
705
                  'digestS3Object': key_name,
706
                  'previousDigestSignature': '...',
707
                  'digestStartTime': (end_date - timedelta(hours=1)).strftime(
708
                      DATE_FORMAT),
709
                  'digestEndTime': end_date.strftime(DATE_FORMAT)}
710
        digest_provider = mock.Mock()
711
        digest_provider.load_digest_keys_in_range.return_value = [
712
            key_name]
713
        digest_provider.fetch_digest.return_value = (digest, key_name)
714
        key_provider = mock.Mock()
715
        public_keys = {'a': {'Fingerprint': 'a', 'Value': 'a'}}
716
        key_provider.get_public_keys.return_value = public_keys
717
        digest_validator = mock.Mock()
718
        traverser = DigestTraverser(
719
            digest_provider=digest_provider, starting_bucket='1',
720
            starting_prefix='baz', public_key_provider=key_provider,
721
            digest_validator=digest_validator)
722
        digest_iter = traverser.traverse(start_date, end_date)
723
        self.assertEqual(digest, next(digest_iter))
724
        digest_validator.validate.assert_called_with(
725
            '1', key_name, public_keys['a']['Value'], digest, key_name)
726

727
    def test_ensures_digest_from_same_location_as_json_contents(self):
728
        start_date = START_DATE
729
        end_date = END_DATE
730
        callback, collected = collecting_callback()
731
        key_name = end_date.strftime(DATE_FORMAT) + '.json.gz'
732
        digest = {'digestPublicKeyFingerprint': 'a',
733
                  'digestS3Bucket': 'not_same',
734
                  'digestS3Object': key_name,
735
                  'digestEndTime': end_date.strftime(DATE_FORMAT)}
736
        digest_provider = mock.Mock()
737
        digest_provider.load_digest_keys_in_range.return_value = [key_name]
738
        digest_provider.fetch_digest.return_value = (digest, key_name)
739
        key_provider = mock.Mock()
740
        digest_validator = mock.Mock()
741
        traverser = DigestTraverser(
742
            digest_provider=digest_provider, starting_bucket='1',
743
            starting_prefix='baz', public_key_provider=key_provider,
744
            digest_validator=digest_validator, on_invalid=callback)
745
        digest_iter = traverser.traverse(start_date, end_date)
746
        self.assertIsNone(next(digest_iter, None))
747
        self.assertEqual(1, len(collected))
748
        self.assertEqual(
749
            'Digest file\ts3://1/%s\tINVALID: invalid format' % key_name,
750
            collected[0]['message'])
751

752
    def test_loads_digests_in_range(self):
753
        start_date = START_DATE
754
        end_date = START_DATE + timedelta(hours=5)
755
        key_provider, digest_provider, validator = create_scenario(
756
            ['gap', 'link', 'link', 'link'])
757
        traverser = DigestTraverser(
758
            digest_provider=digest_provider, starting_bucket='1',
759
            starting_prefix='baz', public_key_provider=key_provider,
760
            digest_validator=validator)
761
        collected = list(traverser.traverse(start_date, end_date))
762
        self.assertEqual(1, key_provider.get_public_keys.call_count)
763
        self.assertEqual(
764
            1, len(digest_provider.calls['load_digest_keys_in_range']))
765
        self.assertEqual(4, len(digest_provider.calls['fetch_digest']))
766
        self.assertEqual(4, len(collected))
767

768
    def test_invokes_cb_and_continues_when_missing(self):
769
        start_date = START_DATE
770
        end_date = END_DATE
771
        key_provider, digest_provider, validator = create_scenario(
772
            ['gap', 'link', 'missing', 'link'])
773
        on_missing, missing_calls = collecting_callback()
774
        traverser = DigestTraverser(
775
            digest_provider=digest_provider, starting_bucket='1',
776
            starting_prefix='baz', public_key_provider=key_provider,
777
            digest_validator=validator, on_missing=on_missing)
778
        collected = list(traverser.traverse(start_date, end_date))
779
        self.assertEqual(3, len(collected))
780
        self.assertEqual(1, key_provider.get_public_keys.call_count)
781
        self.assertEqual(1, len(missing_calls))
782
        # Ensure the keys were provided in the correct order.
783
        self.assertIn('bucket', missing_calls[0])
784
        self.assertIn('next_end_date', missing_calls[0])
785
        # Ensure the keys were provided in the correct order.
786
        self.assertEqual(digest_provider.digests[1],
787
                         missing_calls[0]['next_key'])
788
        self.assertEqual(digest_provider.digests[2],
789
                         missing_calls[0]['last_key'])
790
        # Ensure the provider was called correctly
791
        self.assertEqual(1, key_provider.get_public_keys.call_count)
792
        self.assertEqual(
793
            1, len(digest_provider.calls['load_digest_keys_in_range']))
794
        self.assertEqual(4, len(digest_provider.calls['fetch_digest']))
795

796
    def test_invokes_cb_and_continues_when_invalid(self):
797
        start_date = START_DATE
798
        end_date = END_DATE
799
        key_provider, digest_provider, validator = create_scenario(
800
            ['gap', 'link', 'invalid', 'link', 'invalid'])
801
        on_invalid, invalid_calls = collecting_callback()
802
        traverser = DigestTraverser(
803
            digest_provider=digest_provider, starting_bucket='1',
804
            starting_prefix='baz', public_key_provider=key_provider,
805
            digest_validator=validator, on_invalid=on_invalid)
806
        collected = list(traverser.traverse(start_date, end_date))
807
        self.assertEqual(3, len(collected))
808
        self.assertEqual(1, key_provider.get_public_keys.call_count)
809
        self.assertEqual(2, len(invalid_calls))
810
        # Ensure it was invoked with all the kwargs we expected.
811
        self.assertIn('bucket', invalid_calls[0])
812
        self.assertIn('next_end_date', invalid_calls[0])
813
        # Ensure the keys were provided in the correct order.
814
        self.assertEqual(digest_provider.digests[4],
815
                         invalid_calls[0]['last_key'])
816
        self.assertEqual(digest_provider.digests[3],
817
                         invalid_calls[0]['next_key'])
818
        self.assertEqual(digest_provider.digests[2],
819
                         invalid_calls[1]['last_key'])
820
        self.assertEqual(digest_provider.digests[1],
821
                         invalid_calls[1]['next_key'])
822
        # Ensure the provider was called correctly
823
        self.assertEqual(1, key_provider.get_public_keys.call_count)
824
        self.assertEqual(
825
            1, len(digest_provider.calls['load_digest_keys_in_range']))
826
        self.assertEqual(5, len(digest_provider.calls['fetch_digest']))
827

828
    def test_invokes_cb_and_continues_when_gap(self):
829
        start_date = START_DATE
830
        end_date = END_DATE
831
        key_provider, digest_provider, validator = create_scenario(
832
            ['gap', 'link', 'gap', 'gap'])
833
        on_gap, gap_calls = collecting_callback()
834
        traverser = DigestTraverser(
835
            digest_provider=digest_provider, starting_bucket='1',
836
            starting_prefix='baz', public_key_provider=key_provider,
837
            digest_validator=validator, on_gap=on_gap)
838
        collected = list(traverser.traverse(start_date, end_date))
839
        self.assertEqual(4, len(collected))
840
        self.assertEqual(1, key_provider.get_public_keys.call_count)
841
        self.assertEqual(2, len(gap_calls))
842
        # Ensure it was invoked with all the kwargs we expected.
843
        self.assertIn('bucket', gap_calls[0])
844
        self.assertIn('next_key', gap_calls[0])
845
        self.assertIn('next_end_date', gap_calls[0])
846
        self.assertIn('last_key', gap_calls[0])
847
        self.assertIn('last_start_date', gap_calls[0])
848
        # Ensure the keys were provided in the correct order.
849
        self.assertEqual(digest_provider.digests[3], gap_calls[0]['last_key'])
850
        self.assertEqual(digest_provider.digests[2], gap_calls[0]['next_key'])
851
        self.assertEqual(digest_provider.digests[2], gap_calls[1]['last_key'])
852
        self.assertEqual(digest_provider.digests[1], gap_calls[1]['next_key'])
853
        # Ensure the provider was called correctly
854
        self.assertEqual(1, key_provider.get_public_keys.call_count)
855
        self.assertEqual(
856
            1, len(digest_provider.calls['load_digest_keys_in_range']))
857
        self.assertEqual(4, len(digest_provider.calls['fetch_digest']))
858

859
    def test_reloads_objects_on_bucket_change(self):
860
        start_date = START_DATE
861
        end_date = END_DATE
862
        key_provider, digest_provider, validator = create_scenario(
863
            ['gap', 'link', 'bucket_change', 'link'])
864
        traverser = DigestTraverser(
865
            digest_provider=digest_provider, starting_bucket='1',
866
            starting_prefix='baz', public_key_provider=key_provider,
867
            digest_validator=validator)
868
        collected = list(traverser.traverse(start_date, end_date))
869
        self.assertEqual(4, len(collected))
870
        self.assertEqual(1, key_provider.get_public_keys.call_count)
871
        # Ensure the provider was called correctly
872
        self.assertEqual(1, key_provider.get_public_keys.call_count)
873
        self.assertEqual(
874
            2, len(digest_provider.calls['load_digest_keys_in_range']))
875
        self.assertEqual(['1', '1', '2', '2'],
876
                         [c['digestS3Bucket'] for c in collected])
877

878
    def test_does_not_hard_fail_on_invalid_signature(self):
879
        start_date = START_DATE
880
        end_date = END_DATE
881
        end_timestamp = end_date.strftime(DATE_FORMAT) + '.json.gz'
882
        digest = {'digestPublicKeyFingerprint': 'a',
883
                  'digestS3Bucket': '1',
884
                  'digestS3Object': end_timestamp,
885
                  'previousDigestSignature': '...',
886
                  'digestStartTime': (end_date - timedelta(hours=1)).strftime(
887
                      DATE_FORMAT),
888
                  'digestEndTime': end_timestamp,
889
                  '_signature': '123'}
890
        digest_provider = mock.Mock()
891
        digest_provider.load_digest_keys_in_range.return_value = [
892
            end_timestamp]
893
        digest_provider.fetch_digest.return_value = (digest, end_timestamp)
894
        key_provider = mock.Mock()
895
        public_keys = {'a': {'Fingerprint': 'a', 'Value': 'a'}}
896
        key_provider.get_public_keys.return_value = public_keys
897
        digest_validator = Sha256RSADigestValidator()
898
        on_invalid, calls = collecting_callback()
899
        traverser = DigestTraverser(
900
            digest_provider=digest_provider, starting_bucket='1',
901
            starting_prefix='baz', public_key_provider=key_provider,
902
            digest_validator=digest_validator, on_invalid=on_invalid)
903
        digest_iter = traverser.traverse(start_date, end_date)
904
        next(digest_iter, None)
905
        self.assertIn(
906
            'Digest file\ts3://1/%s\tINVALID: ' % end_timestamp,
907
            calls[0]['message'])
908

909

910
class TestCloudTrailCommand(BaseAWSCommandParamsTest):
911
    def test_s3_client_created_lazily(self):
912
        session = mock.Mock()
913
        command = CloudTrailValidateLogs(session)
914
        parsed_globals = mock.Mock(region=None, verify_ssl=None, endpoint_url=None)
915
        command.setup_services(parsed_globals)
916
        create_client_calls = session.create_client.call_args_list
917
        self.assertEqual(
918
            create_client_calls,
919
            [
920
                mock.call('organizations', verify=None, region_name=None),
921
                mock.call('cloudtrail', verify=None, region_name=None)
922
            ]
923
        )
924

925
    def test_endpoint_url_is_used_for_cloudtrail(self):
926
        endpoint_url = 'https://mycloudtrail.aws.amazon.com/'
927
        session = mock.Mock()
928
        command = CloudTrailValidateLogs(session)
929
        parsed_globals = mock.Mock(region='foo', verify_ssl=None,
930
                              endpoint_url=endpoint_url)
931
        command.setup_services(parsed_globals)
932
        create_client_calls = session.create_client.call_args_list
933
        self.assertEqual(
934
            create_client_calls,
935
            [
936
                mock.call('organizations', verify=None, region_name='foo'),
937
                # Here we should inject the endpoint_url only for cloudtrail.
938
                mock.call('cloudtrail', verify=None, region_name='foo',
939
                     endpoint_url=endpoint_url)
940
            ]
941
        )
942

943
    def test_initializes_args(self):
944
        session = mock.Mock()
945
        command = CloudTrailValidateLogs(session)
946
        start_date = START_DATE.strftime(DATE_FORMAT)
947
        args = Namespace(trail_arn='abc', verbose=True,
948
                         start_time=start_date, s3_bucket='bucket',
949
                         s3_prefix='prefix', end_time=None, account_id=None)
950
        command.handle_args(args)
951
        self.assertEqual('abc', command.trail_arn)
952
        self.assertEqual(True, command.is_verbose)
953
        self.assertEqual('bucket', command.s3_bucket)
954
        self.assertEqual('prefix', command.s3_prefix)
955
        self.assertEqual(start_date, command.start_time.strftime(DATE_FORMAT))
956
        self.assertIsNotNone(command.end_time)
957
        self.assertGreater(command.end_time, command.start_time)
958

959

960
class TestS3ClientProvider(BaseAWSCommandParamsTest):
961
    def test_creates_clients_for_buckets_in_us_east_1(self):
962
        session = mock.Mock()
963
        s3_client = mock.Mock()
964
        session.create_client.return_value = s3_client
965
        s3_client.get_bucket_location.return_value = {'LocationConstraint': ''}
966
        provider = S3ClientProvider(session)
967
        created_client = provider.get_client('foo')
968
        self.assertEqual(s3_client, created_client)
969
        create_client_calls = session.create_client.call_args_list
970
        self.assertEqual(create_client_calls, [mock.call('s3', region_name='us-east-1')])
971
        self.assertEqual(1, s3_client.get_bucket_location.call_count)
972

973
    def test_creates_clients_for_buckets_outside_us_east_1(self):
974
        session = mock.Mock()
975
        s3_client = mock.Mock()
976
        session.create_client.return_value = s3_client
977
        s3_client.get_bucket_location.return_value = {
978
            'LocationConstraint': 'us-west-2'}
979
        provider = S3ClientProvider(session, 'us-west-1')
980
        created_client = provider.get_client('foo')
981
        self.assertEqual(s3_client, created_client)
982
        create_client_calls = session.create_client.call_args_list
983
        self.assertEqual(create_client_calls, [
984
            mock.call('s3', region_name='us-west-1'),
985
            mock.call('s3', region_name='us-west-2')
986
        ])
987
        self.assertEqual(1, s3_client.get_bucket_location.call_count)
988

989
    def test_caches_previously_loaded_bucket_regions(self):
990
        session = mock.Mock()
991
        s3_client = mock.Mock()
992
        session.create_client.return_value = s3_client
993
        s3_client.get_bucket_location.return_value = {'LocationConstraint': ''}
994
        provider = S3ClientProvider(session)
995
        provider.get_client('foo')
996
        self.assertEqual(1, s3_client.get_bucket_location.call_count)
997
        provider.get_client('foo')
998
        self.assertEqual(1, s3_client.get_bucket_location.call_count)
999
        provider.get_client('bar')
1000
        self.assertEqual(2, s3_client.get_bucket_location.call_count)
1001
        provider.get_client('bar')
1002
        self.assertEqual(2, s3_client.get_bucket_location.call_count)
1003

1004
    def test_caches_previously_loaded_clients(self):
1005
        session = mock.Mock()
1006
        s3_client = mock.Mock()
1007
        session.create_client.return_value = s3_client
1008
        s3_client.get_bucket_location.return_value = {'LocationConstraint': ''}
1009
        provider = S3ClientProvider(session)
1010
        client = provider.get_client('foo')
1011
        self.assertEqual(1, session.create_client.call_count)
1012
        self.assertEqual(client, provider.get_client('foo'))
1013
        self.assertEqual(1, session.create_client.call_count)
1014

1015
    def test_removes_cli_error_events(self):
1016
        # We should also remove the error handler for S3.
1017
        # This can be removed once the client switchover is done.
1018
        session = mock.Mock()
1019
        s3_client = mock.Mock()
1020
        session.create_client.return_value = s3_client
1021
        s3_client.get_bucket_location.return_value = {'LocationConstraint': ''}
1022
        provider = S3ClientProvider(session)
1023
        client = provider.get_client('foo')
1024

1025
Product

Resources

Company
