# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License 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.
"""Unit tests for Superset"""
import json
import unittest

import mock

from superset import app, db, security_manager
from superset.connectors.connector_registry import ConnectorRegistry
from superset.connectors.druid.models import DruidDatasource
from superset.connectors.sqla.models import SqlaTable
from superset.models import core as models
from .base_tests import SupersetTestCase

ROLE_TABLES_PERM_DATA = {
    'role_name': 'override_me',
    'database': [{
        'datasource_type': 'table',
        'name': 'main',
        'schema': [{
            'name': '',
            'datasources': ['birth_names'],
        }],
    }],
}

ROLE_ALL_PERM_DATA = {
    'role_name': 'override_me',
    'database': [{
        'datasource_type': 'table',
        'name': 'main',
        'schema': [{
            'name': '',
            'datasources': ['birth_names'],
        }],
    }, {
        'datasource_type': 'druid',
        'name': 'druid_test',
        'schema': [{
            'name': '',
            'datasources': ['druid_ds_1', 'druid_ds_2'],
        }],
    },
    ],
}

EXTEND_ROLE_REQUEST = (
    '/superset/approve?datasource_type={}&datasource_id={}&'
    'created_by={}&role_to_extend={}')
GRANT_ROLE_REQUEST = (
    '/superset/approve?datasource_type={}&datasource_id={}&'
    'created_by={}&role_to_grant={}')
TEST_ROLE_1 = 'test_role1'
TEST_ROLE_2 = 'test_role2'
DB_ACCESS_ROLE = 'db_access_role'
SCHEMA_ACCESS_ROLE = 'schema_access_role'


def create_access_request(session, ds_type, ds_name, role_name, user_name):
    ds_class = ConnectorRegistry.sources[ds_type]
    # TODO: generalize datasource names
    if ds_type == 'table':
        ds = session.query(ds_class).filter(
            ds_class.table_name == ds_name).first()
    else:
        ds = session.query(ds_class).filter(
            ds_class.datasource_name == ds_name).first()
    ds_perm_view = security_manager.find_permission_view_menu(
        'datasource_access', ds.perm)
    security_manager.add_permission_role(
        security_manager.find_role(role_name), ds_perm_view)
    access_request = models.DatasourceAccessRequest(
        datasource_id=ds.id,
        datasource_type=ds_type,
        created_by_fk=security_manager.find_user(username=user_name).id,
    )
    session.add(access_request)
    session.commit()
    return access_request


class RequestAccessTests(SupersetTestCase):

    @classmethod
    def setUpClass(cls):
        security_manager.add_role('override_me')
        security_manager.add_role(TEST_ROLE_1)
        security_manager.add_role(TEST_ROLE_2)
        security_manager.add_role(DB_ACCESS_ROLE)
        security_manager.add_role(SCHEMA_ACCESS_ROLE)
        db.session.commit()

    @classmethod
    def tearDownClass(cls):
        override_me = security_manager.find_role('override_me')
        db.session.delete(override_me)
        db.session.delete(security_manager.find_role(TEST_ROLE_1))
        db.session.delete(security_manager.find_role(TEST_ROLE_2))
        db.session.delete(security_manager.find_role(DB_ACCESS_ROLE))
        db.session.delete(security_manager.find_role(SCHEMA_ACCESS_ROLE))
        db.session.commit()

    def setUp(self):
        self.login('admin')

    def tearDown(self):
        self.logout()
        override_me = security_manager.find_role('override_me')
        override_me.permissions = []
        db.session.commit()
        db.session.close()

    def test_override_role_permissions_is_admin_only(self):
        self.logout()
        self.login('alpha')
        response = self.client.post(
            '/superset/override_role_permissions/',
            data=json.dumps(ROLE_TABLES_PERM_DATA),
            content_type='application/json',
            follow_redirects=True)
        self.assertNotEquals(405, response.status_code)

    def test_override_role_permissions_1_table(self):
        response = self.client.post(
            '/superset/override_role_permissions/',
            data=json.dumps(ROLE_TABLES_PERM_DATA),
            content_type='application/json')
        self.assertEquals(201, response.status_code)

        updated_override_me = security_manager.find_role('override_me')
        self.assertEquals(1, len(updated_override_me.permissions))
        birth_names = self.get_table_by_name('birth_names')
        self.assertEquals(
            birth_names.perm,
            updated_override_me.permissions[0].view_menu.name)
        self.assertEquals(
            'datasource_access',
            updated_override_me.permissions[0].permission.name)

    def test_override_role_permissions_druid_and_table(self):
        response = self.client.post(
            '/superset/override_role_permissions/',
            data=json.dumps(ROLE_ALL_PERM_DATA),
            content_type='application/json')
        self.assertEquals(201, response.status_code)

        updated_role = security_manager.find_role('override_me')
        perms = sorted(
            updated_role.permissions, key=lambda p: p.view_menu.name)
        druid_ds_1 = self.get_druid_ds_by_name('druid_ds_1')
        self.assertEquals(druid_ds_1.perm, perms[0].view_menu.name)
        self.assertEquals('datasource_access', perms[0].permission.name)

        druid_ds_2 = self.get_druid_ds_by_name('druid_ds_2')
        self.assertEquals(druid_ds_2.perm, perms[1].view_menu.name)
        self.assertEquals(
            'datasource_access', updated_role.permissions[1].permission.name)

        birth_names = self.get_table_by_name('birth_names')
        self.assertEquals(birth_names.perm, perms[2].view_menu.name)
        self.assertEquals(
            'datasource_access', updated_role.permissions[2].permission.name)
        self.assertEquals(3, len(perms))

    def test_override_role_permissions_drops_absent_perms(self):
        override_me = security_manager.find_role('override_me')
        override_me.permissions.append(
            security_manager.find_permission_view_menu(
                view_menu_name=self.get_table_by_name('energy_usage').perm,
                permission_name='datasource_access'),
        )
        db.session.flush()

        response = self.client.post(
            '/superset/override_role_permissions/',
            data=json.dumps(ROLE_TABLES_PERM_DATA),
            content_type='application/json')
        self.assertEquals(201, response.status_code)
        updated_override_me = security_manager.find_role('override_me')
        self.assertEquals(1, len(updated_override_me.permissions))
        birth_names = self.get_table_by_name('birth_names')
        self.assertEquals(
            birth_names.perm,
            updated_override_me.permissions[0].view_menu.name)
        self.assertEquals(
            'datasource_access',
            updated_override_me.permissions[0].permission.name)

    def test_clean_requests_after_role_extend(self):
        session = db.session

        # Case 1. Gamma and gamma2 requested test_role1 on energy_usage access
        # Gamma already has role test_role1
        # Extend test_role1 with energy_usage access for gamma2
        # Check if access request for gamma at energy_usage was deleted

        # gamma2 and gamma request table_role on energy usage
        if app.config.get('ENABLE_ACCESS_REQUEST'):
            access_request1 = create_access_request(
                session, 'table', 'random_time_series', TEST_ROLE_1, 'gamma2')
            ds_1_id = access_request1.datasource_id
            create_access_request(
                session, 'table', 'random_time_series', TEST_ROLE_1, 'gamma')
            access_requests = self.get_access_requests('gamma', 'table', ds_1_id)
            self.assertTrue(access_requests)
            # gamma gets test_role1
            self.get_resp(GRANT_ROLE_REQUEST.format(
                'table', ds_1_id, 'gamma', TEST_ROLE_1))
            # extend test_role1 with access on energy usage
            self.client.get(EXTEND_ROLE_REQUEST.format(
                'table', ds_1_id, 'gamma2', TEST_ROLE_1))
            access_requests = self.get_access_requests('gamma', 'table', ds_1_id)
            self.assertFalse(access_requests)

            gamma_user = security_manager.find_user(username='gamma')
            gamma_user.roles.remove(security_manager.find_role('test_role1'))

    def test_clean_requests_after_alpha_grant(self):
        session = db.session

        # Case 2. Two access requests from gamma and gamma2
        # Gamma becomes alpha, gamma2 gets granted
        # Check if request by gamma has been deleted

        access_request1 = create_access_request(
            session, 'table', 'birth_names', TEST_ROLE_1, 'gamma')
        create_access_request(
            session, 'table', 'birth_names', TEST_ROLE_2, 'gamma2')
        ds_1_id = access_request1.datasource_id
        # gamma becomes alpha
        alpha_role = security_manager.find_role('Alpha')
        gamma_user = security_manager.find_user(username='gamma')
        gamma_user.roles.append(alpha_role)
        session.commit()
        access_requests = self.get_access_requests('gamma', 'table', ds_1_id)
        self.assertTrue(access_requests)
        self.client.get(EXTEND_ROLE_REQUEST.format(
            'table', ds_1_id, 'gamma2', TEST_ROLE_2))
        access_requests = self.get_access_requests('gamma', 'table', ds_1_id)
        self.assertFalse(access_requests)

        gamma_user = security_manager.find_user(username='gamma')
        gamma_user.roles.remove(security_manager.find_role('Alpha'))
        session.commit()

    def test_clean_requests_after_db_grant(self):
        session = db.session

        # Case 3. Two access requests from gamma and gamma2
        # Gamma gets database access, gamma2 access request granted
        # Check if request by gamma has been deleted

        gamma_user = security_manager.find_user(username='gamma')
        access_request1 = create_access_request(
            session, 'table', 'energy_usage', TEST_ROLE_1, 'gamma')
        create_access_request(
            session, 'table', 'energy_usage', TEST_ROLE_2, 'gamma2')
        ds_1_id = access_request1.datasource_id
        # gamma gets granted database access
        database = session.query(models.Database).first()

        security_manager.merge_perm('database_access', database.perm)
        ds_perm_view = security_manager.find_permission_view_menu(
            'database_access', database.perm)
        security_manager.add_permission_role(
            security_manager.find_role(DB_ACCESS_ROLE), ds_perm_view)
        gamma_user.roles.append(security_manager.find_role(DB_ACCESS_ROLE))
        session.commit()
        access_requests = self.get_access_requests('gamma', 'table', ds_1_id)
        self.assertTrue(access_requests)
        # gamma2 request gets fulfilled
        self.client.get(EXTEND_ROLE_REQUEST.format(
            'table', ds_1_id, 'gamma2', TEST_ROLE_2))
        access_requests = self.get_access_requests('gamma', 'table', ds_1_id)

        self.assertFalse(access_requests)
        gamma_user = security_manager.find_user(username='gamma')
        gamma_user.roles.remove(security_manager.find_role(DB_ACCESS_ROLE))
        session.commit()

    def test_clean_requests_after_schema_grant(self):
        session = db.session

        # Case 4. Two access requests from gamma and gamma2
        # Gamma gets schema access, gamma2 access request granted
        # Check if request by gamma has been deleted

        gamma_user = security_manager.find_user(username='gamma')
        access_request1 = create_access_request(
            session, 'table', 'wb_health_population', TEST_ROLE_1, 'gamma')
        create_access_request(
            session, 'table', 'wb_health_population', TEST_ROLE_2, 'gamma2')
        ds_1_id = access_request1.datasource_id
        ds = session.query(SqlaTable).filter_by(
            table_name='wb_health_population').first()

        ds.schema = 'temp_schema'
        security_manager.merge_perm('schema_access', ds.schema_perm)
        schema_perm_view = security_manager.find_permission_view_menu(
            'schema_access', ds.schema_perm)
        security_manager.add_permission_role(
            security_manager.find_role(SCHEMA_ACCESS_ROLE), schema_perm_view)
        gamma_user.roles.append(security_manager.find_role(SCHEMA_ACCESS_ROLE))
        session.commit()
        # gamma2 request gets fulfilled
        self.client.get(EXTEND_ROLE_REQUEST.format(
            'table', ds_1_id, 'gamma2', TEST_ROLE_2))
        access_requests = self.get_access_requests('gamma', 'table', ds_1_id)
        self.assertFalse(access_requests)
        gamma_user = security_manager.find_user(username='gamma')
        gamma_user.roles.remove(security_manager.find_role(SCHEMA_ACCESS_ROLE))

        ds = session.query(SqlaTable).filter_by(
            table_name='wb_health_population').first()
        ds.schema = None

        session.commit()

    @mock.patch('superset.utils.core.send_MIME_email')
    def test_approve(self, mock_send_mime):
        if app.config.get('ENABLE_ACCESS_REQUEST'):
            session = db.session
            TEST_ROLE_NAME = 'table_role'
            security_manager.add_role(TEST_ROLE_NAME)

            # Case 1. Grant new role to the user.

            access_request1 = create_access_request(
                session, 'table', 'unicode_test', TEST_ROLE_NAME, 'gamma')
            ds_1_id = access_request1.datasource_id
            self.get_resp(GRANT_ROLE_REQUEST.format(
                'table', ds_1_id, 'gamma', TEST_ROLE_NAME))
            # Test email content.
            self.assertTrue(mock_send_mime.called)
            call_args = mock_send_mime.call_args[0]
            self.assertEqual([security_manager.find_user(username='gamma').email,
                              security_manager.find_user(username='admin').email],
                             call_args[1])
            self.assertEqual(
                '[Superset] Access to the datasource {} was granted'.format(
                    self.get_table(ds_1_id).full_name), call_args[2]['Subject'])
            self.assertIn(TEST_ROLE_NAME, call_args[2].as_string())
            self.assertIn('unicode_test', call_args[2].as_string())

            access_requests = self.get_access_requests('gamma', 'table', ds_1_id)
            # request was removed
            self.assertFalse(access_requests)
            # user was granted table_role
            user_roles = [r.name for r in security_manager.find_user('gamma').roles]
            self.assertIn(TEST_ROLE_NAME, user_roles)

            # Case 2. Extend the role to have access to the table

            access_request2 = create_access_request(
                session, 'table', 'energy_usage', TEST_ROLE_NAME, 'gamma')
            ds_2_id = access_request2.datasource_id
            energy_usage_perm = access_request2.datasource.perm

            self.client.get(EXTEND_ROLE_REQUEST.format(
                'table', access_request2.datasource_id, 'gamma', TEST_ROLE_NAME))
            access_requests = self.get_access_requests('gamma', 'table', ds_2_id)

            # Test email content.
            self.assertTrue(mock_send_mime.called)
            call_args = mock_send_mime.call_args[0]
            self.assertEqual([security_manager.find_user(username='gamma').email,
                              security_manager.find_user(username='admin').email],
                             call_args[1])
            self.assertEqual(
                '[Superset] Access to the datasource {} was granted'.format(
                    self.get_table(ds_2_id).full_name), call_args[2]['Subject'])
            self.assertIn(TEST_ROLE_NAME, call_args[2].as_string())
            self.assertIn('energy_usage', call_args[2].as_string())

            # request was removed
            self.assertFalse(access_requests)
            # table_role was extended to grant access to the energy_usage table/
            perm_view = security_manager.find_permission_view_menu(
                'datasource_access', energy_usage_perm)
            TEST_ROLE = security_manager.find_role(TEST_ROLE_NAME)
            self.assertIn(perm_view, TEST_ROLE.permissions)

            # Case 3. Grant new role to the user to access the druid datasource.

            security_manager.add_role('druid_role')
            access_request3 = create_access_request(
                session, 'druid', 'druid_ds_1', 'druid_role', 'gamma')
            self.get_resp(GRANT_ROLE_REQUEST.format(
                'druid', access_request3.datasource_id, 'gamma', 'druid_role'))

            # user was granted table_role
            user_roles = [r.name for r in security_manager.find_user('gamma').roles]
            self.assertIn('druid_role', user_roles)

            # Case 4. Extend the role to have access to the druid datasource

            access_request4 = create_access_request(
                session, 'druid', 'druid_ds_2', 'druid_role', 'gamma')
            druid_ds_2_perm = access_request4.datasource.perm

            self.client.get(EXTEND_ROLE_REQUEST.format(
                'druid', access_request4.datasource_id, 'gamma', 'druid_role'))
            # druid_role was extended to grant access to the druid_access_ds_2
            druid_role = security_manager.find_role('druid_role')
            perm_view = security_manager.find_permission_view_menu(
                'datasource_access', druid_ds_2_perm)
            self.assertIn(perm_view, druid_role.permissions)

            # cleanup
            gamma_user = security_manager.find_user(username='gamma')
            gamma_user.roles.remove(security_manager.find_role('druid_role'))
            gamma_user.roles.remove(security_manager.find_role(TEST_ROLE_NAME))
            session.delete(security_manager.find_role('druid_role'))
            session.delete(security_manager.find_role(TEST_ROLE_NAME))
            session.commit()

    def test_request_access(self):
        if app.config.get('ENABLE_ACCESS_REQUEST'):
            session = db.session
            self.logout()
            self.login(username='gamma')
            gamma_user = security_manager.find_user(username='gamma')
            security_manager.add_role('dummy_role')
            gamma_user.roles.append(security_manager.find_role('dummy_role'))
            session.commit()

            ACCESS_REQUEST = (
                '/superset/request_access?'
                'datasource_type={}&'
                'datasource_id={}&'
                'action={}&')
            ROLE_GRANT_LINK = (
                '<a href="/superset/approve?datasource_type={}&datasource_id={}&'
                'created_by={}&role_to_grant={}">Grant {} Role</a>')

            # Request table access, there are no roles have this table.

            table1 = session.query(SqlaTable).filter_by(
                table_name='random_time_series').first()
            table_1_id = table1.id

            # request access to the table
            resp = self.get_resp(
                ACCESS_REQUEST.format('table', table_1_id, 'go'))
            assert 'Access was requested' in resp
            access_request1 = self.get_access_requests('gamma', 'table', table_1_id)
            assert access_request1 is not None

            # Request access, roles exist that contains the table.
            # add table to the existing roles
            table3 = session.query(SqlaTable).filter_by(
                table_name='energy_usage').first()
            table_3_id = table3.id
            table3_perm = table3.perm

            security_manager.add_role('energy_usage_role')
            alpha_role = security_manager.find_role('Alpha')
            security_manager.add_permission_role(
                alpha_role,
                security_manager.find_permission_view_menu(
                    'datasource_access', table3_perm))
            security_manager.add_permission_role(
                security_manager.find_role('energy_usage_role'),
                security_manager.find_permission_view_menu(
                    'datasource_access', table3_perm))
            session.commit()

            self.get_resp(
                ACCESS_REQUEST.format('table', table_3_id, 'go'))
            access_request3 = self.get_access_requests('gamma', 'table', table_3_id)
            approve_link_3 = ROLE_GRANT_LINK.format(
                'table', table_3_id, 'gamma', 'energy_usage_role',
                'energy_usage_role')
            self.assertEqual(access_request3.roles_with_datasource,
                             '<ul><li>{}</li></ul>'.format(approve_link_3))

            # Request druid access, there are no roles have this table.
            druid_ds_4 = session.query(DruidDatasource).filter_by(
                datasource_name='druid_ds_1').first()
            druid_ds_4_id = druid_ds_4.id

            # request access to the table
            self.get_resp(ACCESS_REQUEST.format('druid', druid_ds_4_id, 'go'))
            access_request4 = self.get_access_requests('gamma', 'druid', druid_ds_4_id)

            self.assertEqual(
                access_request4.roles_with_datasource,
                '<ul></ul>'.format(access_request4.id))

            # Case 5. Roles exist that contains the druid datasource.
            # add druid ds to the existing roles
            druid_ds_5 = session.query(DruidDatasource).filter_by(
                datasource_name='druid_ds_2').first()
            druid_ds_5_id = druid_ds_5.id
            druid_ds_5_perm = druid_ds_5.perm

            druid_ds_2_role = security_manager.add_role('druid_ds_2_role')
            admin_role = security_manager.find_role('Admin')
            security_manager.add_permission_role(
                admin_role,
                security_manager.find_permission_view_menu(
                    'datasource_access', druid_ds_5_perm))
            security_manager.add_permission_role(
                druid_ds_2_role,
                security_manager.find_permission_view_menu(
                    'datasource_access', druid_ds_5_perm))
            session.commit()

            self.get_resp(ACCESS_REQUEST.format('druid', druid_ds_5_id, 'go'))
            access_request5 = self.get_access_requests(
                'gamma', 'druid', druid_ds_5_id)
            approve_link_5 = ROLE_GRANT_LINK.format(
                'druid', druid_ds_5_id, 'gamma', 'druid_ds_2_role',
                'druid_ds_2_role')
            self.assertEqual(access_request5.roles_with_datasource,
                             '<ul><li>{}</li></ul>'.format(approve_link_5))

            # cleanup
            gamma_user = security_manager.find_user(username='gamma')
            gamma_user.roles.remove(security_manager.find_role('dummy_role'))
            session.commit()


if __name__ == '__main__':
    unittest.main()