<reponame>sloria/flask-resty
import json

from marshmallow import fields, Schema
import pytest
from sqlalchemy import Column, Integer, String

from flask_resty import Api, GenericModelView, HasAnyCredentialsAuthorization
from flask_resty.testing import assert_response

# -----------------------------------------------------------------------------

try:
    from flask_resty import JwtAuthentication, JwkSetAuthentication
except ImportError:
    pytestmark = pytest.mark.skip(reason="JWT support not installed")

# -----------------------------------------------------------------------------


class AbstractTestJwt(object):
    @pytest.yield_fixture
    def models(self, db):
        class Widget(db.Model):
            __tablename__ = 'widgets'

            id = Column(Integer, primary_key=True)
            owner_id = Column(String)

        db.create_all()

        yield {
            'widget': Widget,
        }

        db.drop_all()

    @pytest.fixture
    def schemas(self):
        class WidgetSchema(Schema):
            id = fields.Integer(as_string=True)
            owner_id = fields.String()

        return {
            'widget': WidgetSchema(),
        }

    @pytest.fixture
    def auth(self, app):
        raise NotImplementedError()

    @pytest.fixture(autouse=True)
    def routes(self, app, models, schemas, auth):
        class WidgetListView(GenericModelView):
            model = models['widget']
            schema = schemas['widget']

            authentication = auth['authentication']
            authorization = auth['authorization']

            def get(self):
                return self.list()

        api = Api(app)
        api.add_resource('/widgets', WidgetListView)

    @pytest.fixture(autouse=True)
    def data(self, db, models):
        db.session.add_all((
            models['widget'](owner_id='foo'),
            models['widget'](owner_id='bar'),
        ))
        db.session.commit()

    @pytest.fixture
    def token(self):
        raise NotImplementedError()

    @pytest.fixture
    def invalid_token(self, request):
        raise NotImplementedError()

    def test_header(self, client, token):
        response = client.get(
            '/widgets',
            headers={
                'Authorization': 'Bearer {}'.format(token),
            },
        )

        assert_response(response, 200, [
            {
                'id': '1',
                'owner_id': 'foo',
            },
        ])

    def test_arg(self, client, token):
        response = client.get('/widgets?id_token={}'.format(token))
        assert_response(response, 200, [
            {
                'id': '1',
                'owner_id': 'foo',
            },
        ])

    def test_error_unauthenticated(self, client):
        response = client.get('/widgets')
        assert_response(response, 401, [{
            'code': 'invalid_credentials.missing',
        }])

    def test_error_invalid_authorization(self, client):
        response = client.get(
            '/widgets',
            headers={
                'Authorization': 'foo',
            },
        )
        assert_response(response, 401, [{
            'code': 'invalid_authorization',
        }])

    def test_error_invalid_authorization_scheme(self, client):
        response = client.get(
            '/widgets',
            headers={
                'Authorization': 'foo bar',
            },
        )
        assert_response(response, 401, [{
            'code': 'invalid_authorization.scheme',
        }])

    def test_error_invalid_token(self, client, invalid_token):
        response = client.get(
            '/widgets',
            headers={
                'Authorization': 'Bearer {}'.format(invalid_token),
            },
        )
        assert_response(response, 401, [{
            'code': 'invalid_token',
        }])


# -----------------------------------------------------------------------------


class TestJwt(AbstractTestJwt):
    @pytest.fixture
    def auth(self, app):
        app.config.update({
            'RESTY_JWT_DECODE_KEY': 'secret',
            'RESTY_JWT_DECODE_ALGORITHMS': ['HS256'],
        })
        authentication = JwtAuthentication(issuer='resty')

        class UserAuthorization(HasAnyCredentialsAuthorization):
            def filter_query(self, query, view):
                return query.filter_by(
                    owner_id=self.get_request_credentials()['sub'],
                )

        return {
            'authentication': authentication,
            'authorization': UserAuthorization(),
        }

    @pytest.fixture
    def token(self):
        return '<KEY>'  # noqa: E501

    @pytest.fixture(
        params=(
            'foo',
            '<KEY>',  # noqa: E501
            '<KEY>',  # noqa: E501
        ),
        ids=(
            'malformed',
            'key_mismatch',
            'iss_mismatch',
        ),
    )
    def invalid_token(self, request):
        return request.param


class TestJwkSet(AbstractTestJwt):
    @pytest.fixture(
        params=(
            'tests/fixtures/testkey_rsa_pub.json',
            'tests/fixtures/testkey_rsa_cert.json',
        ),
        ids=('public key', 'cert'),
    )
    def jwk_set(self, request):
        with open(request.param, 'r') as rsa_pub_file:
            return json.load(rsa_pub_file)

    @pytest.fixture()
    def auth(self, app, jwk_set):
        app.config.update({
            'RESTY_JWT_DECODE_JWK_SET': jwk_set,
            'RESTY_JWT_DECODE_ALGORITHMS': ['RS256'],
        })

        authentication = JwkSetAuthentication(issuer='resty')

        class UserAuthorization(HasAnyCredentialsAuthorization):
            def filter_query(self, query, view):
                return query.filter_by(
                    owner_id=self.get_request_credentials()['sub'],
                )

        return {
            'authentication': authentication,
            'authorization': UserAuthorization(),
        }

    @pytest.fixture
    def token(self):
        return '<KEY>'  # noqa: E501

    @pytest.fixture(
        params=(
            'foo',
            '<KEY>',  # noqa: E501
            '<KEY>',  # noqa: E501
            '<KEY>',  # noqa: E501
            '<KEY>',  # noqa: E501
            '<KEY>',  # noqa: E501
        ),
        ids=(
            'malformed',
            'key_mismatch',
            'iss_mismatch',
            'kid_mismatch',
            'wrong_algorithm',
            'missing_alg',
        ),
    )
    def invalid_token(self, request):
        return request.param
