# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
import re
from contextlib import contextmanager
from nose.tools import assert_equal, assert_true, assert_raises
from tests import unittest, mock, BaseSessionTest, ClientHTTPStubber

from botocore import exceptions
from botocore.exceptions import (
    UnsupportedS3ControlArnError,
    UnsupportedS3ControlConfigurationError,
    InvalidHostLabelError,
    ParamValidationError,
)
from botocore.session import Session
from botocore.compat import urlsplit
from botocore.config import Config
from botocore.awsrequest import AWSResponse


ACCESSPOINT_ARN_TEST_CASES = [
    # Outpost accesspoint arn test cases
    {
        'arn': (
            'arn:aws:s3-outposts:us-west-2:123456789012:'
            'outpost:op-01234567890123456:accesspoint:myaccesspoint'
        ),
        'region': 'us-west-2',
        'config': {},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-west-2.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-east-1:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'region': 'us-west-2',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-east-1.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-east-1:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'region': 'us-west-2',
        'config': {'s3': {'use_arn_region': False}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws-cn:s3-outposts:cn-north-1:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'region': 'us-west-2',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws-us-gov:s3-outposts:us-gov-east-1:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'region': 'us-gov-east-1',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-gov-east-1.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws-us-gov:s3-outposts:us-gov-east-1:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'region': 'us-gov-east-1-fips',
        'config': {'s3': {'use_arn_region': False}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    #{
    #    'arn': 'arn:aws-us-gov:s3-outposts:fips-us-gov-east-1:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
    #    'region': 'fips-us-gov-east-1',
    #    'config': {'s3': {'use_arn_region': True}},
    #    'assertions': {
    #        'exception': 'UnsupportedS3ArnError',
    #    }
    #},
    {
        'arn': 'arn:aws-us-gov:s3-outposts:us-gov-east-1:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'region': 'us-gov-east-1-fips',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-gov-east-1.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'config': {'s3': {'use_dualstack_endpoint': True}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:op-01234567890123456:accesspoint:myaccesspoint',
        'config': {'s3': {'use_accelerate_endpoint': True}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost',
        'assertions': {
            'exception': 'UnsupportedS3ControlArnError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:op-01234567890123456',
        'assertions': {
            'exception': 'UnsupportedS3ControlArnError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:myaccesspoint',
        'assertions': {
            'exception': 'UnsupportedS3ControlArnError',
        }
    },
]


BUCKET_ARN_TEST_CASES = [
    # Outpost bucket arn test cases
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-west-2',
        'config': {},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-west-2.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-east-1:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-west-2',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-east-1.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-east-1:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-west-2',
        'config': {'s3': {'use_arn_region': False}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws-cn:s3-outposts:cn-north-1:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-west-2',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws-us-gov:s3-outposts:us-gov-east-1:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-gov-east-1',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-gov-east-1.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws-us-gov:s3-outposts:us-gov-east-1:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-gov-east-1-fips',
        'config': {'s3': {'use_arn_region': False}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws-us-gov:s3-outposts:fips-us-gov-east-1:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'fips-us-gov-east-1',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws-us-gov:s3-outposts:us-gov-east-1:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-gov-east-1-fips',
        'config': {'s3': {'use_arn_region': True}},
        'assertions': {
            'signing_name': 's3-outposts',
            'netloc': 's3-outposts.us-gov-east-1.amazonaws.com',
            'headers': {
                'x-amz-outpost-id': 'op-01234567890123456',
                'x-amz-account-id': '123456789012',
            },
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:op-01234567890123456:bucket:mybucket',
        'region': 'us-west-2',
        'config': {'s3': {'use_dualstack_endpoint': True}},
        'assertions': {
            'exception': 'UnsupportedS3ControlConfigurationError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost',
        'assertions': {
            'exception': 'UnsupportedS3ControlArnError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:op-01234567890123456',
        'assertions': {
            'exception': 'UnsupportedS3ControlArnError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:bucket',
        'assertions': {
            'exception': 'UnsupportedS3ControlArnError',
        }
    },
    {
        'arn': 'arn:aws:s3-outposts:us-west-2:123456789012:outpost:op-01234567890123456:bucket',
        'assertions': {
            'exception': 'UnsupportedS3ControlArnError',
        }
    },
]


V4_AUTH_REGEX = re.compile(
    r'AWS4-HMAC-SHA256 Credential=\w+/\d+/'
    r'(?P<region>[a-z0-9-]+)/(?P<name>[a-z0-9-]+)/'
)


def _assert_signing_name(stubber, expected_name):
    request = stubber.requests[0]
    auth_header = request.headers['Authorization'].decode('utf-8')
    actual_name = V4_AUTH_REGEX.match(auth_header).group('name')
    assert_equal(expected_name, actual_name)


def _assert_netloc(stubber, expected_netloc):
    request = stubber.requests[0]
    url_parts = urlsplit(request.url)
    assert_equal(expected_netloc, url_parts.netloc)


def _assert_header(stubber, key, value):
    request = stubber.requests[0]
    assert_true(key in request.headers)
    actual_value = request.headers[key]
    if isinstance(actual_value, bytes):
        actual_value = actual_value.decode('utf-8')
    assert_equal(value, actual_value)


def _assert_headers(stubber, headers):
    for key, value in headers.items():
        _assert_header(stubber, key, value)


def _bootstrap_session():
    session = Session()
    session.set_credentials('access_key', 'secret_key')
    return session


def _bootstrap_client(session, region, **kwargs):
    client = session.create_client('s3control', region, **kwargs)
    stubber = ClientHTTPStubber(client)
    return client, stubber


def _bootstrap_test_case_client(session, test_case):
    region = test_case.get('region', 'us-west-2')
    config = test_case.get('config', {})
    config = Config(**config)
    return _bootstrap_client(session, region, config=config)


def test_accesspoint_arn_redirection():
    session = _bootstrap_session()
    for test_case in ACCESSPOINT_ARN_TEST_CASES:
        client, stubber = _bootstrap_test_case_client(session, test_case)
        yield _test_accesspoint_arn, test_case, client, stubber


def _test_accesspoint_arn(test_case, client, stubber):
    with _assert_test_case(test_case, client, stubber):
        client.get_access_point_policy(Name=test_case['arn'])


def test_bucket_arn_redirection():
    session = _bootstrap_session()
    for test_case in BUCKET_ARN_TEST_CASES:
        client, stubber = _bootstrap_test_case_client(session, test_case)
        yield _test_bucket_arn, test_case, client, stubber


def _test_bucket_arn(test_case, client, stubber):
    with _assert_test_case(test_case, client, stubber):
        client.get_bucket(Bucket=test_case['arn'])


@contextmanager
def _assert_test_case(test_case, client, stubber):
    stubber.add_response()
    assertions = test_case['assertions']
    exception_raised = None
    try:
        with stubber:
            yield
    except Exception as e:
        if 'exception' not in assertions:
            raise
        exception_raised = e
    if 'exception' in assertions:
        exception_cls = getattr(exceptions, assertions['exception'])
        if exception_raised is None:
            raise RuntimeError(
                'Expected exception "%s" was not raised' % exception_cls
            )
        error_msg = (
            'Expected exception "%s", got "%s"'
        ) % (exception_cls, type(exception_raised))
        assert isinstance(exception_raised, exception_cls), error_msg
    else:
        assert_equal(len(stubber.requests), 1)
        if 'signing_name' in assertions:
            _assert_signing_name(stubber, assertions['signing_name'])
        if 'headers' in assertions:
            _assert_headers(stubber, assertions['headers'])
        if 'netloc' in assertions:
            _assert_netloc(stubber, assertions['netloc'])


class TestS3ControlRedirection(unittest.TestCase):
    def setUp(self):
        self.session = _bootstrap_session()
        self.region = 'us-west-2'
        self._bootstrap_client()

    def _bootstrap_client(self, **kwargs):
        client, stubber = _bootstrap_client(
            self.session, self.region, **kwargs
        )
        self.client = client
        self.stubber = stubber

    def test_outpost_id_redirection_dualstack(self):
        config = Config(s3={'use_dualstack_endpoint': True})
        self._bootstrap_client(config=config)
        with self.assertRaises(UnsupportedS3ControlConfigurationError):
            self.client.create_bucket(Bucket='foo', OutpostId='op-123')

    def test_outpost_id_redirection_create_bucket(self):
        self.stubber.add_response()
        with self.stubber:
            self.client.create_bucket(Bucket='foo', OutpostId='op-123')
        _assert_netloc(self.stubber, 's3-outposts.us-west-2.amazonaws.com')
        _assert_header(self.stubber, 'x-amz-outpost-id', 'op-123')

    def test_outpost_id_redirection_list_regional_buckets(self):
        self.stubber.add_response()
        with self.stubber:
            self.client.list_regional_buckets(
                OutpostId='op-123',
                AccountId='1234',
            )
        _assert_netloc(self.stubber, 's3-outposts.us-west-2.amazonaws.com')
        _assert_header(self.stubber, 'x-amz-outpost-id', 'op-123')

    def test_outpost_redirection_custom_endpoint(self):
        self._bootstrap_client(endpoint_url='https://outpost.foo.com/')
        self.stubber.add_response()
        with self.stubber:
            self.client.create_bucket(Bucket='foo', OutpostId='op-123')
        _assert_netloc(self.stubber, 'outpost.foo.com')
        _assert_header(self.stubber, 'x-amz-outpost-id', 'op-123')

    def test_normal_ap_request_has_correct_endpoint(self):
        self.stubber.add_response()
        with self.stubber:
            self.client.get_access_point_policy(Name='MyAp', AccountId='1234')
        _assert_netloc(self.stubber, '1234.s3-control.us-west-2.amazonaws.com')

    def test_normal_ap_request_custom_endpoint(self):
        self._bootstrap_client(endpoint_url='https://example.com/')
        self.stubber.add_response()
        with self.stubber:
            self.client.get_access_point_policy(Name='MyAp', AccountId='1234')
        _assert_netloc(self.stubber, '1234.example.com')

    def test_normal_bucket_request_has_correct_endpoint(self):
        self.stubber.add_response()
        with self.stubber:
            self.client.create_bucket(Bucket='foo')
        _assert_netloc(self.stubber, 's3-control.us-west-2.amazonaws.com')

    def test_arn_account_id_mismatch(self):
        arn = (
            'arn:aws:s3-outposts:us-west-2:123456789012:outpost:'
            'op-01234567890123456:accesspoint:myaccesspoint'
        )
        with self.assertRaises(UnsupportedS3ControlArnError):
            self.client.get_access_point_policy(Name=arn, AccountId='1234')

    def test_create_access_point_does_not_expand_name(self):
        arn = (
            'arn:aws:s3-outposts:us-west-2:123456789012:outpost:'
            'op-01234567890123456:accesspoint:myaccesspoint'
        )
        # AccountId will not be backfilled by the arn expansion and thus
        # fail parameter validation
        with self.assertRaises(ParamValidationError):
            self.client.create_access_point(Name=arn, Bucket='foo')

    def test_arn_invalid_host_label(self):
        config = Config(s3={'use_arn_region': True})
        self._bootstrap_client(config=config)
        arn = (
            'arn:aws:s3-outposts:us-we$t-2:123456789012:outpost:'
            'op-01234567890123456:accesspoint:myaccesspoint'
        )
        with self.assertRaises(InvalidHostLabelError):
            self.client.get_access_point_policy(Name=arn)

    def test_unknown_arn_format(self):
        arn = 'arn:aws:foo:us-west-2:123456789012:bar:myresource'
        with self.assertRaises(UnsupportedS3ControlArnError):
            self.client.get_access_point_policy(Name=arn)
