# Copyright 2018 RethinkDB
#
# 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.
#
# This file incorporates work covered by the following copyright:
# Copyright 2010-2016 RethinkDB, all rights reserved.

import six
import base64
import binascii
import hashlib
import hmac
import struct
import sys
import threading

from random import SystemRandom
from rethinkdb import ql2_pb2
from rethinkdb.errors import ReqlAuthError, ReqlDriverError
from rethinkdb.helpers import decode_utf8, chain_to_bytes
from rethinkdb.logger import default_logger


try:
    xrange
except NameError:
    xrange = range


def compare_digest(digest_a, digest_b):
    if sys.version_info[0] == 3:
        def xor_bytes(digest_a, digest_b):
            return digest_a ^ digest_b
    else:
        def xor_bytes(digest_a, digest_b, _ord=ord):
            return _ord(digest_a) ^ _ord(digest_b)

    left = None
    right = digest_b
    if len(digest_a) == len(digest_b):
        left = digest_a
        result = 0
    if len(digest_a) != len(digest_b):
        left = digest_b
        result = 1

    for l, r in zip(left, right):
        result |= xor_bytes(l, r)

    return result == 0


def pbkdf2_hmac(hash_name, password, salt, iterations):
    if hash_name != 'sha256':
        raise AssertionError('Hash name {hash_name} is not equal with "sha256"'.format(hash_name=hash_name))

    def from_bytes(value, hexlify=binascii.hexlify, int=int):
        return int(hexlify(value), 16)

    def to_bytes(value, unhexlify=binascii.unhexlify):
        try:
            return unhexlify(bytes('%064x' % value, 'ascii'))
        except TypeError:
            return unhexlify(bytes('%064x' % value))

    cache_key = (password, salt, iterations)

    cache_result = HandshakeV1_0.PBKDF2_CACHE.get(cache_key)

    if cache_result is not None:
        return cache_result

    mac = hmac.new(password, None, hashlib.sha256)

    def digest(msg, mac=mac):
        mac_copy = mac.copy()
        mac_copy.update(msg)
        return mac_copy.digest()

    t = digest(salt + b'\x00\x00\x00\x01')
    u = from_bytes(t)
    for c in xrange(iterations - 1):
        t = digest(t)
        u ^= from_bytes(t)

    u = to_bytes(u)
    HandshakeV1_0.PBKDF2_CACHE.set(cache_key, u)
    return u


class LocalThreadCache(threading.local):
    def __init__(self):
        self._cache = dict()

    def set(self, key, val):
        self._cache[key] = val

    def get(self, key):
        return self._cache.get(key)


class HandshakeV1_0(object):
    """
    RethinkDB client drivers are responsible for serializing queries, sending them to the server using the
    ReQL wire protocol, and receiving responses from the server and returning them to the calling application.

    The client sends the protocol version, authentication method, and authentication as a null-terminated JSON
    response. RethinkDB currently supports only one authentication method, SCRAM-SHA-256, as specified in IETF
    RFC 7677 and RFC 5802. The RFC is followed with the exception of error handling (RethinkDB uses its own
    higher level error reporting rather than the e= field). RethinkDB does not support channel binding and clients
    should not request this. The value of "authentication" is the "client-first-message" specified in RFC 5802
    (the channel binding flag, optional SASL authorization identity, username (n=), and random nonce (r=).

    More info: https://rethinkdb.com/docs/writing-drivers/
    """

    VERSION = ql2_pb2.VersionDummy.Version.V1_0
    PROTOCOL = ql2_pb2.VersionDummy.Protocol.JSON
    PBKDF2_CACHE = LocalThreadCache()

    def __init__(self, json_decoder, json_encoder, host, port, username, password):
        """
        TODO:
        """

        self._json_decoder = json_decoder
        self._json_encoder = json_encoder
        self._host = host
        self._port = port
        self._username = username.encode('utf-8').replace(b'=', b'=3D').replace(b',', b'=2C')

        self._password = six.b(password)

        self._compare_digest = self._get_compare_digest()
        self._pbkdf2_hmac = self._get_pbkdf2_hmac()

        self._protocol_version = 0
        self._random_nonce = None
        self._first_client_message = None
        self._server_signature = None
        self._state = 0

    @staticmethod
    def _get_compare_digest():
        """
        Get the compare_digest function from hashlib if package contains it, else get
        our own function. Please note that hashlib contains this function only for
        Python 2.7.7+ and 3.3+.
        """

        return getattr(hmac, 'compare_digest', compare_digest)

    @staticmethod
    def _get_pbkdf2_hmac():
        """
        Get the pbkdf2_hmac function from hashlib if package contains it, else get
        our own function. Please note that hashlib contains this function only for
        Python 2.7.8+ and 3.4+.
        """

        return getattr(hashlib, 'pbkdf2_hmac', pbkdf2_hmac)

    @staticmethod
    def _get_authentication_and_first_client_message(response):
        """
        Get the first client message and the authentication related data from the
        response provided by RethinkDB.

        :param response: Response dict from the database
        :return: None
        """

        first_client_message = response['authentication'].encode('ascii')
        authentication = dict(x.split(b'=', 1) for x in first_client_message.split(b','))
        return first_client_message, authentication

    def _next_state(self):
        """
        Increase the state counter.
        """

        self._state += 1

    def _decode_json_response(self, response, with_utf8=False):
        """
        Get decoded json response from response.

        :param response: Response from the database
        :param with_utf8: UTF-8 decode response before json decoding
        :raises: ReqlDriverError | ReqlAuthError
        :return: Json decoded response of the original response
        """

        if with_utf8:
            response = decode_utf8(response)

        json_response = self._json_decoder.decode(response)

        if not json_response.get('success'):
            if 10 <= json_response['error_code'] <= 20:
                raise ReqlAuthError(json_response['error'], self._host, self._port)

            raise ReqlDriverError(json_response['error'])

        return json_response

    def _init_connection(self, response):
        """
        Prepare initial connection message. We send the version as well as the initial
        JSON as an optimization.

        :param response: Response from the database
        :raises: ReqlDriverError
        :return: Initial message which will be sent to the DB
        """

        if response is not None:
            raise ReqlDriverError('Unexpected response')

        self._random_nonce = base64.standard_b64encode(bytes(bytearray(
            SystemRandom().getrandbits(8) for i in range(18)
        )))

        self._first_client_message = chain_to_bytes('n=', self._username, ',r=', self._random_nonce)

        initial_message = chain_to_bytes(
            struct.pack('<L', self.VERSION),
            self._json_encoder.encode({
                'protocol_version': self._protocol_version,
                'authentication_method': 'SCRAM-SHA-256',
                'authentication': chain_to_bytes('n,,', self._first_client_message).decode('ascii')
            }).encode('utf-8'),
            b'\0'
        )

        self._next_state()
        return initial_message

    def _read_response(self, response):
        """
        Read response of the server. Due to we've already sent the initial JSON, and only support a single
        protocol version at the moment thus we simply read the next response and return an empty string as a
        message.

        :param response: Response from the database
        :raises: ReqlDriverError | ReqlAuthError
        :return: An empty string
        """


        json_response = self._decode_json_response(response)
        min_protocol_version = json_response['min_protocol_version']
        max_protocol_version = json_response['max_protocol_version']

        if not min_protocol_version <= self._protocol_version <= max_protocol_version:
            raise ReqlDriverError('Unsupported protocol version {version}, expected between {min} and {max}'.format(
                version=self._protocol_version,
                min=min_protocol_version,
                max=max_protocol_version
            ))

        self._next_state()
        return ''

    def _prepare_auth_request(self, response):
        """
        Put tohether the authentication request based on the response of the database.

        :param response: Response from the database
        :raises: ReqlDriverError | ReqlAuthError
        :return: An empty string
        """

        json_response = self._decode_json_response(response, with_utf8=True)
        first_client_message, authentication = self._get_authentication_and_first_client_message(json_response)

        random_nonce = authentication[b'r']
        if not random_nonce.startswith(self._random_nonce):
            raise ReqlAuthError('Invalid nonce from server', self._host, self._port)

        salted_password = self._pbkdf2_hmac(
            'sha256',
            self._password,
            base64.standard_b64decode(authentication[b's']),
            int(authentication[b'i'])
        )

        message_without_proof = chain_to_bytes('c=biws,r=', random_nonce)
        auth_message = b','.join((
            self._first_client_message,
            first_client_message,
            message_without_proof
        ))

        self._server_signature = hmac.new(
            hmac.new(salted_password, b'Server Key', hashlib.sha256).digest(),
            auth_message,
            hashlib.sha256
        ).digest()

        client_key = hmac.new(salted_password, b'Client Key', hashlib.sha256).digest()
        client_signature = hmac.new(hashlib.sha256(client_key).digest(), auth_message, hashlib.sha256).digest()
        client_proof = struct.pack('32B', *(l ^ random_nonce for l, random_nonce in zip(
            struct.unpack('32B', client_key),
            struct.unpack('32B', client_signature)
        )))

        authentication_request = chain_to_bytes(
            self._json_encoder.encode({
                'authentication': chain_to_bytes(
                    message_without_proof, ',p=', base64.standard_b64encode(client_proof)
                ).decode('ascii')
            }),
            b'\0'
        )

        self._next_state()
        return authentication_request

    def _read_auth_response(self, response):
        """
        Read the authentication request's response sent by the database
        and validate the server signature which was returned.

        :param response: Response from the database
        :raises: ReqlDriverError | ReqlAuthError
        :return: None
        """

        json_response = self._decode_json_response(response, with_utf8=True)

        first_client_message, authentication = self._get_authentication_and_first_client_message(json_response)
        server_signature = base64.standard_b64decode(authentication[b'v'])

        if not self._compare_digest(server_signature, self._server_signature):
            raise ReqlAuthError('Invalid server signature', self._host, self._port)

        self._next_state()

    def reset(self):
        self._random_nonce = None
        self._first_client_message = None
        self._server_signature = None
        self._state = 0

    def next_message(self, response):
        if response is not None:
            response = response.decode('utf-8')

        if self._state == 0:
            return self._init_connection(response)

        elif self._state == 1:
            return self._read_response(response)

        elif self._state == 2:
            return self._prepare_auth_request(response)

        elif self._state == 3:
            return self._read_auth_response(response)

        raise ReqlDriverError('Unexpected handshake state')
