# Copyright 2019 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 hashlib
import json
import os
import re

from botocore.context import start_as_current_context

from awscli.testutils import mock
from tests.functional.sso import BaseSSOTest


class TestLoginCommand(BaseSSOTest):
    _TIMESTAMP_FORMAT_PATTERN = re.compile(
        r'\d\d\d\d-\d\d-\d\dT\d\d:\d\d:\d\dZ'
    )

    def assert_cache_contains_registration(
        self, start_url, session_name, scopes, expected_client_id
    ):
        cached_files = os.listdir(self.token_cache_dir)

        cached_registration_filename = self._get_cached_registration_filename(
            start_url, session_name, scopes
        )

        self.assertIn(cached_registration_filename, cached_files)
        self.assertEqual(
            self._get_cached_response(cached_registration_filename)[
                'clientId'
            ],
            expected_client_id,
        )

    def assert_cache_contains_token(
        self,
        start_url,
        expected_token,
        session_name=None,
    ):
        cached_files = os.listdir(self.token_cache_dir)
        # The registration and cached access token
        self.assertEqual(len(cached_files), 2)
        cached_token_filename = self._get_cached_token_filename(
            start_url,
            session_name,
        )
        self.assertIn(cached_token_filename, cached_files)
        self.assertEqual(
            self._get_cached_response(cached_token_filename)['accessToken'],
            expected_token,
        )

    def _get_cached_registration_filename(
        self, start_url, session_name, scopes
    ):
        args = {
            'tool': 'botocore',
            'startUrl': start_url,
            'region': self.sso_region,
            'scopes': scopes,
            'session_name': session_name,
        }
        cache_args = json.dumps(args, sort_keys=True).encode('utf-8')
        return hashlib.sha1(cache_args).hexdigest() + '.json'

    def _get_cached_token_filename(self, start_url, session_name):
        to_hash = start_url
        if session_name:
            to_hash = session_name
        return hashlib.sha1(to_hash.encode('utf-8')).hexdigest() + '.json'

    def _get_cached_response(self, token_filename):
        token_path = os.path.join(self.token_cache_dir, token_filename)
        with open(token_path) as f:
            cached_response = json.loads(f.read())
            return cached_response

    def assert_cache_token_expiration_time_format_is_correct(self):
        token_filename = self._get_cached_token_filename(self.start_url, None)
        token_path = os.path.join(self.token_cache_dir, token_filename)
        with open(token_path) as f:
            cached_response = json.loads(f.read())
            self.assertIsNotNone(
                self._TIMESTAMP_FORMAT_PATTERN.match(
                    cached_response['expiresAt']
                )
            )

    def test_login_explicit_device(self):
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --use-device-code')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_device_browser_handler_called_with(
            'foo',
            'https://sso.fake/device',
            'https://sso.verify',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url, expected_token=self.access_token
        )

    def test_login_implicit_device(self):
        # This is a legacy profile via setUp, so we expect
        # it to fall back to device flow automatically
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_device_browser_handler_called_with(
            'foo',
            'https://sso.fake/device',
            'https://sso.verify',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url, expected_token=self.access_token
        )

    def test_login_device_no_browser(self):
        self.add_oidc_device_responses(self.access_token)
        stdout, _, _ = self.run_cmd('sso login --use-device-code --no-browser')
        self.assertIn('Browser will not be automatically opened.', stdout)
        self.open_browser_mock.assert_not_called()
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_cache_contains_token(
            start_url=self.start_url, expected_token=self.access_token
        )

    def test_login_auth_no_browser(self):
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)
        stdout, _, _ = self.run_cmd('sso login --no-browser')
        self.assertIn('Browser will not be automatically opened.', stdout)
        self.open_browser_mock.assert_not_called()
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='auth-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            expected_token=self.access_token,
            session_name='test-session',
        )

    def test_login_device_forces_refresh(self):
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --use-device-code')
        # The register response from the first login should have been
        # cached.
        self.add_oidc_device_responses(
            'new.token', include_register_response=False
        )
        self.run_cmd('sso login --use-device-code')
        self.assert_cache_contains_token(
            start_url=self.start_url,
            expected_token='new.token',
        )

    def test_login_auth_forces_refresh(self):
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        # The register response from the first login should have been
        # cached.
        self.add_oidc_auth_code_responses(
            'new.token', include_register_response=False
        )
        self.run_cmd('sso login')
        self.assert_cache_contains_token(
            start_url=self.start_url,
            expected_token='new.token',
            session_name='test-session',
        )

    def test_login_auth_after_device_forces_refresh(self):
        self.set_config_file_content(
            content=self.get_sso_session_config('test-session')
        )
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --use-device-code')
        # The register response from the first login should have been
        # cached.
        self.add_oidc_auth_code_responses('new.token')
        self.run_cmd('sso login')
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='auth-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            expected_token='new.token',
            session_name='test-session',
        )

    def test_login_device_no_sso_configuration(self):
        self.set_config_file_content(content='')
        _, stderr, _ = self.run_cmd(
            'sso login --use-device-code', expected_rc=253
        )
        self.assertIn(
            'Missing the following required SSO configuration', stderr
        )

    def test_login_auth_no_sso_configuration(self):
        self.set_config_file_content(content='')
        _, stderr, _ = self.run_cmd('sso login', expected_rc=253)
        self.assertIn(
            'Missing the following required SSO configuration', stderr
        )

    def test_login_device_minimal_sso_configuration(self):
        content = (
            '[default]\n'
            f'sso_start_url={self.start_url}\n'
            f'sso_region={self.sso_region}\n'
        )
        self.set_config_file_content(content=content)
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_device_browser_handler_called_with(
            'foo',
            'https://sso.fake/device',
            'https://sso.verify',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url, expected_token=self.access_token
        )

    def test_login_device_partially_missing_sso_configuration(self):
        content = f'[default]\nsso_start_url={self.start_url}\n'
        self.set_config_file_content(content=content)
        _, stderr, _ = self.run_cmd(
            'sso login --use-device-code', expected_rc=253
        )
        self.assertIn(
            'Missing the following required SSO configuration', stderr
        )
        self.assertIn('sso_region', stderr)
        self.assertNotIn('sso_start_url', stderr)
        self.assertNotIn('sso_account_id', stderr)
        self.assertNotIn('sso_role_name', stderr)

    def test_token_cache_datetime_format(self):
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --use-device-code')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_cache_contains_token(
            start_url=self.start_url, expected_token=self.access_token
        )
        self.assert_cache_token_expiration_time_format_is_correct()

    def test_login_device_sso_session(self):
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --use-device-code')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_device_browser_handler_called_with(
            'foo',
            'https://sso.fake/device',
            'https://sso.verify',
        )
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='device-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            session_name='test-session',
            expected_token=self.access_token,
        )

    def test_login_auth_sso_session(self):
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_auth_browser_handler_called_with('sso%3Aaccount%3Aaccess')
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='auth-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            session_name='test-session',
            expected_token=self.access_token,
        )

    def test_login_device_sso_with_explicit_sso_session_arg(self):
        content = self.get_sso_session_config(
            'test-session', include_profile=False
        )
        self.set_config_file_content(content=content)
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --sso-session test-session --use-device-code')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_device_browser_handler_called_with(
            'foo',
            'https://sso.fake/device',
            'https://sso.verify',
        )
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='device-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            session_name='test-session',
            expected_token=self.access_token,
        )

    def test_login_auth_sso_with_explicit_sso_session_arg(self):
        content = self.get_sso_session_config(
            'test-session', include_profile=False
        )
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login --sso-session test-session')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_auth_browser_handler_called_with('sso%3Aaccount%3Aaccess')
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='auth-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            session_name='test-session',
            expected_token=self.access_token,
        )

    def test_login_device_sso_session_with_scopes(self):
        self.registration_scopes = ['sso:foo', 'sso:bar']
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --use-device-code')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_device_browser_handler_called_with(
            'foo',
            'https://sso.fake/device',
            'https://sso.verify',
        )
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='device-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            session_name='test-session',
            expected_token=self.access_token,
        )
        operation, params = self.operations_called[0]
        self.assertEqual(operation.name, 'RegisterClient')
        self.assertEqual(params.get('scopes'), self.registration_scopes)

    def test_login_auth_sso_session_with_scopes(self):
        self.registration_scopes = ['sso:foo', 'sso:bar']
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        self.assert_used_expected_sso_region(expected_region=self.sso_region)
        self.assert_auth_browser_handler_called_with('sso%3Afoo%20sso%3Abar')
        self.assert_cache_contains_registration(
            start_url=self.start_url,
            session_name='test-session',
            scopes=self.registration_scopes,
            expected_client_id='auth-client-id',
        )
        self.assert_cache_contains_token(
            start_url=self.start_url,
            session_name='test-session',
            expected_token=self.access_token,
        )
        operation, params = self.operations_called[0]
        self.assertEqual(operation.name, 'RegisterClient')
        self.assertEqual(params.get('scopes'), self.registration_scopes)

    def test_login_sso_session_missing_config(self):
        content = '[default]\n' 'sso_session=test\n' '[sso-session test]\n'
        self.set_config_file_content(content=content)
        _, stderr, _ = self.run_cmd('sso login', expected_rc=253)
        self.assertIn(
            'SSO configuration values: sso_start_url, sso_region', stderr
        )

    def test_login_sso_session_missing(self):
        content = '[default]\n' 'sso_session=test\n'
        self.set_config_file_content(content=content)
        _, stderr, _ = self.run_cmd('sso login', expected_rc=253)
        self.assertIn('sso-session does not exist: "test"', stderr)

    def test_login_auth_sso_no_authorization_code_throws_error(self):
        self.fetcher_mock.return_value.get_auth_code_and_state.return_value = (
            None,
            None,
        )
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)

        _, stderr, _ = self.run_cmd('sso login', expected_rc=255)
        self.assertIn('Failed to retrieve an authorization code.', stderr)

    def test_login_auth_sso_state_mismatch_throws_error(self):
        self.fetcher_mock.return_value.get_auth_code_and_state.return_value = (
            "abc",
            '00000000-0000-0000-0000-000000000001',
        )
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)

        _, stderr, _ = self.run_cmd('sso login', expected_rc=255)
        self.assertIn('State parameter does not match expected value.', stderr)

    def test_login_device_no_extra_user_agent(self):
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login --use-device-code')
        self.assertNotIn(
            'md/sso#auth', self.last_request_dict['headers']['User-Agent']
        )

    def test_login_auth_includes_extra_user_agent(self):
        content = self.get_sso_session_config('test-session')
        self.set_config_file_content(content=content)
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        self.assertIn(
            'md/sso#auth', self.last_request_dict['headers']['User-Agent']
        )

    @mock.patch('botocore.session.Session.create_client')
    def test_login_custom_ca(self, create_client_patch):
        # Profile and mock responses aren't wired up so this fails with 255,
        # but it does get far enough to verify that the internal SSO-OIDC
        # client is created with the custom CA bundle
        self.run_cmd(
            'sso login --ca-bundle /path/to/ca/bundle.pem', expected_rc=255
        )
        create_client_patch.assert_called_once_with(
            'sso-oidc', config=mock.ANY, verify='/path/to/ca/bundle.pem'
        )

    @mock.patch('botocore.session.Session.create_client')
    def test_login_no_verify_ssl(self, create_client_patch):
        # Profile and mock responses aren't wired up so this fails with 255,
        # but it does get far enough to verify that the internal SSO-OIDC
        # client is created with verify=False
        self.run_cmd('sso login --no-verify-ssl', expected_rc=255)
        create_client_patch.assert_called_once_with(
            'sso-oidc',
            config=mock.ANY,
            verify=False,
        )


class TestLoginWithVanityUrl(BaseSSOTest):
    def setUp(self):
        super().setUp()
        self.vanity_url = 'https://aws.mycompany.com'
        self.resolved_url = 'https://ssoins-abc.portal.us-east-1.app.aws'
        self.start_url = self.vanity_url
        self.set_config_file_content(
            self.get_sso_session_config('vanity-session')
        )
        self._resolve_patch.stop()
        self.resolve_patch = mock.patch(
            'awscli.customizations.sso.login.resolve_start_url',
            return_value=(self.resolved_url, 'us-east-1'),
        )
        self.mock_resolve = self.resolve_patch.start()

    def tearDown(self):
        super().tearDown()
        self.resolve_patch.stop()

    def test_login_uses_resolved_region(self):
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        self.assert_used_expected_sso_region('us-east-1')

    def test_login_passes_resolved_url_to_register_client(self):
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        register_call = self.operations_called[0]
        self.assertEqual(register_call[0].name, 'RegisterClient')
        self.assertEqual(register_call[1]['issuerUrl'], self.resolved_url)

    def test_config_rewritten_after_successful_login(self):
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        with open(self.config_file) as f:
            config_content = f.read()
        self.assertIn('sso_region = us-east-1', config_content)

    def test_region_update_message_printed(self):
        self.add_oidc_auth_code_responses(self.access_token)
        stdout, _, _ = self.run_cmd('sso login')
        self.assertIn('SSO region updated to us-east-1', stdout)

    def test_no_region_update_message_when_region_matches(self):
        self.sso_region = 'us-east-1'
        self.set_config_file_content(
            self.get_sso_session_config('vanity-session')
        )
        self.add_oidc_auth_code_responses(self.access_token)
        stdout, _, _ = self.run_cmd('sso login')
        self.assertNotIn('SSO region updated', stdout)

    def test_config_not_rewritten_on_login_failure(self):
        self.parsed_responses = [
            {
                'Error': {
                    'Code': 'InvalidRequestException',
                    'Message': 'Invalid start url provided',
                }
            }
        ]
        self.run_cmd('sso login', expected_rc=254)
        with open(self.config_file) as f:
            config_content = f.read()
        self.assertIn(f'sso_region={self.sso_region}', config_content)

    def test_config_not_rewritten_when_region_matches(self):
        self.sso_region = 'us-east-1'
        self.set_config_file_content(
            self.get_sso_session_config('vanity-session')
        )
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        with open(self.config_file) as f:
            config_content = f.read()
        self.assertIn('sso_region=us-east-1', config_content)

    def test_vanity_url_sets_feature_id_in_user_agent(self):
        self.add_oidc_auth_code_responses(self.access_token)
        with start_as_current_context():
            self.run_cmd('sso login')
        self.assertIn('AM', self.last_request_dict['headers']['User-Agent'])


class TestLoginWithVanityUrlLegacy(BaseSSOTest):
    def setUp(self):
        super().setUp()
        self.vanity_url = 'https://aws.mycompany.com'
        self.resolved_url = 'https://ssoins-abc.portal.us-east-1.app.aws'
        self.start_url = self.vanity_url
        self.set_config_file_content(self.get_legacy_config())
        self._resolve_patch.stop()
        self.resolve_patch = mock.patch(
            'awscli.customizations.sso.login.resolve_start_url',
            return_value=(self.resolved_url, 'us-east-1'),
        )
        self.mock_resolve = self.resolve_patch.start()

    def tearDown(self):
        super().tearDown()
        self.resolve_patch.stop()

    def test_legacy_login_uses_resolved_region(self):
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login')
        self.assert_used_expected_sso_region('us-east-1')

    def test_legacy_config_rewritten_after_successful_login(self):
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login')
        with open(self.config_file) as f:
            config_content = f.read()
        self.assertIn('sso_region = us-east-1', config_content)

    def test_legacy_passes_resolved_url_to_device_authorization(self):
        self.add_oidc_device_responses(self.access_token)
        self.run_cmd('sso login')
        device_auth_call = self.operations_called[1]
        self.assertEqual(device_auth_call[0].name, 'StartDeviceAuthorization')
        self.assertEqual(device_auth_call[1]['startUrl'], self.resolved_url)


class TestLoginWithDirectUrl(BaseSSOTest):
    def setUp(self):
        super().setUp()
        self.start_url = 'https://ssoins-abc.portal.us-west-2.app.aws'
        self.set_config_file_content(
            self.get_sso_session_config('direct-session')
        )
        self._resolve_patch.stop()
        self.resolve_patch = mock.patch(
            'awscli.customizations.sso.login.resolve_start_url',
            return_value=(self.start_url, 'us-west-2'),
        )
        self.mock_resolve = self.resolve_patch.start()

    def tearDown(self):
        super().tearDown()
        self.resolve_patch.stop()

    def test_config_not_rewritten_for_direct_url(self):
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        with open(self.config_file) as f:
            config_content = f.read()
        self.assertIn(f'sso_region={self.sso_region}', config_content)

    def test_login_uses_configured_region(self):
        self.add_oidc_auth_code_responses(self.access_token)
        self.run_cmd('sso login')
        self.assert_used_expected_sso_region('us-west-2')

    def test_direct_url_does_not_set_vanity_feature_id(self):
        self.add_oidc_auth_code_responses(self.access_token)
        with start_as_current_context():
            self.run_cmd('sso login')
        user_agent = self.last_request_dict['headers']['User-Agent']
        self.assertNotIn('AM', user_agent.split('m/')[-1])
