#!/usr/bin/env python
# -*- coding: utf-8 -*-

###############################################################################
#  Copyright Kitware Inc.
#
#  Licensed 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.
###############################################################################

import re

import pytest

from girder.exceptions import AccessException
from girder.models.user import User
from pytest_girder.assertions import assertStatus, assertStatusOk


def testInitializeOtp(user):
    otpUris = User().initializeOtp(user)

    # A URI for TOTP should be returned
    assert re.match(r'', otpUris['totpUri'])

    # OTP should not be enabled yet, since it's not finalized
    assert user['otp']['enabled'] is False
    # TOTP parameters should be generated
    assert 'totp' in user['otp']


def testHasOtp(user):
    assert User().hasOtp(user) is False

    User().initializeOtp(user)

    # OTP is not yet enabled
    assert User().hasOtp(user) is False

    user['otp']['enabled'] = True

    assert User().hasOtp(user) is True


def _tokenFromTotpUri(totpUri, valid=True):
    # Create an external TOTP instance
    from passlib.totp import TOTP
    totp = TOTP.from_uri(totpUri)

    # Generate a valid token
    otpToken = totp.generate().token

    if not valid:
        # Increment the token by 1 to invalidate it
        otpToken = str((int(otpToken) + 1) % int(1e6))

    return otpToken


def testVerifyOtp(user):
    # Enable OTP
    otpUris = User().initializeOtp(user)
    user['otp']['enabled'] = True

    # Generate an invalid token
    otpToken = _tokenFromTotpUri(otpUris['totpUri'], False)

    with pytest.raises(AccessException):
        User().verifyOtp(user, otpToken)

    # Generate a valid token
    otpToken = _tokenFromTotpUri(otpUris['totpUri'])
    # Verify the token, which should succeed without raising an exception
    User().verifyOtp(user, otpToken)

    # Re-verify the same token, which should fail
    with pytest.raises(AccessException):
        User().verifyOtp(user, otpToken)


def testAuthenticateWithOtp(user):
    # Providing an unnecessary token should fail
    with pytest.raises(AccessException):
        User().authenticate('user', 'password', '123456')

    # Enable OTP and save user
    otpUris = User().initializeOtp(user)
    user['otp']['enabled'] = True
    User().save(user)

    # Providing no token should now fail
    with pytest.raises(AccessException):
        User().authenticate('user', 'password')

    # Generate a valid token
    otpToken = _tokenFromTotpUri(otpUris['totpUri'])

    # Authenticate successfully with the valid token
    User().authenticate('user', 'password', otpToken)


def testOtpAPIWorkflow(server, user):
    # Try to finalize OTP before it's been initialized
    resp = server.request(
        path='/user/%s/otp' % user['_id'], method='PUT', user=user,
        additionalHeaders=[('Girder-OTP', '123456')])
    # This should fail cleanly
    assertStatus(resp, 400)
    assert 'not initialized' in resp.json['message']

    # Initialize OTP
    resp = server.request(path='/user/%s/otp' % user['_id'], method='POST', user=user)
    assertStatusOk(resp)
    # Save the URI
    totpUri = resp.json['totpUri']

    # Login without an OTP
    resp = server.request(path='/user/authentication', method='GET', basicAuth='user:password')
    # Since OTP has not been finalized, this should still succeed
    assertStatusOk(resp)

    # Finalize without an OTP
    resp = server.request(
        path='/user/%s/otp' % user['_id'], method='PUT', user=user)
    assertStatus(resp, 400)
    assert 'Girder-OTP' in resp.json['message']

    # Finalize with an invalid OTP
    resp = server.request(
        path='/user/%s/otp' % user['_id'], method='PUT', user=user,
        additionalHeaders=[('Girder-OTP', _tokenFromTotpUri(totpUri, False))])
    assertStatus(resp, 403)
    assert 'validation failed' in resp.json['message']

    # Finalize with a valid OTP
    resp = server.request(
        path='/user/%s/otp' % user['_id'], method='PUT', user=user,
        additionalHeaders=[('Girder-OTP', _tokenFromTotpUri(totpUri))])
    assertStatusOk(resp)

    # Login without an OTP
    resp = server.request(path='/user/authentication', method='GET', basicAuth='user:password')
    assertStatus(resp, 401)
    assert 'Girder-OTP' in resp.json['message']

    # Login with an invalid OTP
    resp = server.request(
        path='/user/authentication', method='GET', basicAuth='user:password',
        additionalHeaders=[('Girder-OTP', _tokenFromTotpUri(totpUri, False))])
    assertStatus(resp, 401)
    assert 'validation failed' in resp.json['message']

    # Login with a valid OTP
    resp = server.request(
        path='/user/authentication', method='GET', basicAuth='user:password',
        additionalHeaders=[('Girder-OTP', _tokenFromTotpUri(totpUri))])
    assertStatusOk(resp)

Add OTP testing