# -*- coding: utf-8 -*-

import os
import re
import sys
import json
import time
import uuid
import hmac
import base64
import hashlib
import inspect
import argparse

import bottle
from pddb import PandasDatabase


# Python 2.x vs 3.x compatibility
def tob(s, enc='utf8'):
    if sys.version_info >= (3, 0, 0) and isinstance(s, str):
        s = s.encode(enc)
    elif sys.version_info < (3, 0, 0) and isinstance(s, unicode):
        s = s.encode(enc)
    return s
def tos(s, enc='utf8'):
    if sys.version_info >= (3, 0, 0) and isinstance(s, bytes):
        s = s.decode(enc)
    elif sys.version_info < (3, 0, 0) and isinstance(s, unicode):
        s = s.encode(enc)
    return s

def _lscmp(a, b):
    ''' Compares two strings in a cryptographically safe way:
        Runtime is not affected by length of common prefix. '''
    return not sum(0 if x==y else 1 for x, y in zip(a, b)) and len(a) == len(b)

def data_encode(data, key, digestmod=None):
    ''' Encode and sign a pickle-able object. Return a (byte) string '''
    digestmod = digestmod or hashlib.sha256
    data, key = base64.b64encode(tob(data)), tob(key)
    sig = base64.b64encode(hmac.new(key, data, digestmod=digestmod).digest())
    return tob('!') + sig + tob('?') + data

def data_decode(data, key, digestmod=None):
    ''' Verify and decode an encoded string. Return an object or None.'''
    digestmod = digestmod or hashlib.sha256
    data, key = tob(data), tob(key)
    if data_is_encoded(data):
        sig, data = data.split(tob('?'), 1)
        if _lscmp(sig[1:], base64.b64encode(hmac.new(key, data, digestmod=digestmod).digest())):
            return tos(base64.b64decode(data))
    return None

def data_is_encoded(data):
    ''' Return True if the argument looks like a encoded cookie.'''
    datab = tob(data)
    return bool(datab.startswith(tob('!')) and tob('?') in datab)

class BottleShip(bottle.Bottle):
    '''
    Subclass of bottle.Bottle.

    Each instance of this class corresponds to a single, distinct web application. See
    documentation on bottle.Bottle for additional information about this class and its methods.

    Options
    -------
    pddb : PandasDatabase
        Use the given PandasDatabase object instead of instantiating a new one.
    whitelist_cond : dict
        Conditions to match a new user against. When provided, each item in this dictionary will
        be compared for equality against the `user_info` dictionary passed as an argument to the
        `BottleShip.register()` function.
    blacklist_cond : dict
        Conditions to avoid matching a new user against. When provided, each item in the
        `user_info` dictionary passed as an argument to the `BottleShip.register()` function will
        be compared for inequality against this dictionary.
    allowed_security : array-like (default `["plaintext", "plaintext+ipaddr", "hmac", "hmac+ipaddr"]`)
        Type of security allowed for user authentication. Possible values are:
        ["plaintext", "plaintext+ipaddr", "hmac", "hmac+ipaddr"]
    token_lifetime_seconds : int (default `3600`)
        Number of seconds after which a user token expires.
    max_tokens_per_user : int (default `5`)
        Maximum number of tokens allowed per user. After the limit is reached, new token
        associations will trigger immediate expiration of the oldest token assigned to the user.
    catchall : boolean (default `True`)
        See bottle.Bottle for more information.
    autojson : boolean (default `True`)
        See bottle.Bottle for more information.
    debug : boolean (default `True`)
        When true, prints information about internal operations to the console.
    '''

    def __init__(self, pddb=None, whitelist_cond=None, blacklist_cond=None, allowed_security=None,
                 token_lifetime_seconds=3600, max_tokens_per_user=5, catchall=True, autojson=True,
                 debug=False):

        super(BottleShip, self).__init__()
        self.pddb = pddb or PandasDatabase('bottleship_db', debug=debug)
        self.allowed_security = allowed_security or ('plaintext', 'plaintext+ipaddr', 'hmac', 'hmac+ipaddr')

        table_names = ('bottleship_users', 'bottleship_tokens')
        self.pddb.load(table_names)

        self._rsakey_public = None
        self._rsakey_private = None
        self._key_store = dict()
        self._token_lifetime_seconds = token_lifetime_seconds
        self._whitelist_cond = whitelist_cond or dict()
        self._blacklist_cond = blacklist_cond or dict()
        self._debug = debug

        # If rsa is allowed, generate key pair now
        if any(['rsa' in sec for sec in self.allowed_security]):
            
            # Only import the cryptography module if necessary to avoid it as a dependency
            try:
                from cryptography.hazmat.backends import default_backend
                from cryptography.hazmat.primitives import hashes
                from cryptography.hazmat.primitives.asymmetric import rsa
                from cryptography.hazmat.primitives.asymmetric import padding
                from cryptography.hazmat.primitives.serialization import load_pem_public_key
            except ImportError:
                self._print('RSA is not supported in this system because the cryptography module '
                            'is not installed. To install it, you can run `$ pip install '
                            'cryptography`. RSA encryption has been disabled for this instance.')
                self.allowed_security = [sec for sec in self.allowed_security if 'rsa' not in sec]
            
            self._rsakey_private = rsa.generate_private_key(
                public_exponent=65537, key_size=2048, backend=default_backend())    
            self._rsakey_public = self._rsakey_private.public_key()        
            
    def _print(self, txt):
        if self._debug:
            print(txt)

    def _check_user(self, user_record):
        ''' Test the user record against our whitelist and blacklist '''
        cond_keys = set(list(self._blacklist_cond.keys()) + list(self._whitelist_cond.keys()))

        for cond in cond_keys:
            #self.pddb._regex_type TODO

            if cond in self._whitelist_cond:
                whitelist_cond = self._whitelist_cond[cond]
                if not isinstance(whitelist_cond, (list, tuple)):
                    whitelist_cond = (whitelist_cond,)
                if any([user_record.get(cond) != wc for wc in whitelist_cond]):
                    return False

            if cond in self._blacklist_cond:
                blacklist_cond = self._blacklist_cond[cond]
                if not isinstance(blacklist_cond, (list, tuple)):
                    blacklist_cond = (blacklist_cond,)
                if any([user_record.get(cond) == bc for bc in blacklist_cond]):
                    return False

        return True

    def _gen_token(self, username, security_level=None, key=None):
        ''' Generates a new token for a username and stores it in the database '''

        # Generate the appropriate token for the given security level
        security_level = security_level or self.allowed_security[0]
        if security_level in ('plaintext', 'plaintext+ipaddr', 'hmac', 'hmac+ipaddr'):
            token = str(uuid.uuid4()) # TODO
        else:
            raise NotImplementedError(
                'Token could not be generated for requested security level')

        # Set expiration time, username, etc.
        token_record = {
            'Token': token,
            'Expiry': str(time.time() + self._token_lifetime_seconds),
            'Username': username,
            'SecurityLevel': security_level,
            'Key': key or '',
        }
        print(token_record)
        token_record = self.pddb.insert('bottleship_tokens', record=token_record, astype='dict')
        return token_record

    def _dump_user_record(self, security_level, user_record):
        user_record_json = json.dumps(user_record)
        if 'hmac' in security_level:
            user_record_json = data_encode(user_record_json, user_record.get('Key'))
        elif 'rsa' in security_level:
            user_publickey = load_pem_public_key(
                user_record.get('Key'), backend=default_backend())
            if not isinstance(user_publickey, rsa.RSAPublicKey):
                msg = 'Login error: User key could not be loaded as RSA public key.'
                self._print(msg)
                return bottle.HTTPResponse(status=403, body=msg)
            user_record_json = user_publickey.encrypt(
                user_record_json, padding.OAEP(mgf=padding.MGF1(
                    algorithm=hashes.SHA1()), algorithm=hashes.SHA1(), label=None))
        return user_record_json

    def _read_secure_json(self, data):
        '''
        Read secure JSON data in the format {"Data": <secure_json>, "Token": <single-use token>}.
        Returns <data_dict, error_message>
        '''

        err_msg = None
        if data.get('Token') not in self._key_store:
            err_msg = ('Security error: Single-use token could not be found in the key store. '
                       'Please request a new one using the key exchange API.')
            return None, err_msg

        # Retrieve and delete single-use token
        security_level, user_key, server_key = self._key_store.pop(data.get('Token'))

        data_json = None
        if 'hmac' in security_level:
            if data_is_encoded(data.get('Data')):
                data_json = data_decode(data.get('Data'), server_key)
            else:
                err_msg = ('Security error: Provided data was not encoded but the security level '
                           'requires it.')
                return None, err_msg
        elif 'rsa' in security_level:
            data_json = self._rsakey_private.decrypt(data.get('Data'), padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA1()),
                algorithm=hashes.SHA1(), label=None))
        else:
            err_msg = ('Security error: Provided security level is not supported. Must be one of: '
                        '%r.' % [sec for sec in self.allowed_security if 'hmac' in sec or 'rsa' in sec])
            return None, err_msg

        # Attempt to parse the data into a dictionary and return it
        data_dict = None
        try:
            data_dict = json.loads(tob(data_json).decode('unicode_escape').strip('"'))
            data_dict['Key'] = user_key
        except TypeError:
            err_msg = ('Security error: Provided data was corrupted and could not be parsed as a '
                        'JSON object.')
            return None, err_msg

        return data_dict, err_msg

    def _error_username_password(self, username, password):
        ''' Returns error message if verification fails, None otherwise '''

        # Make sure that username is not empty
        if not username:
            return 'Username error: Parameter "username" cannot be empty.'

        # Make sure that username and password are of the correct type
        if not isinstance(username, str) or not isinstance(password, str):
            return 'Credential error: Parameters "username" and "password" must be of type "str".'

    def key_exchange(self, security_level, user_key):
        '''
        Exchange keys with client.

        Parameters
        ----------
        security_level : str
            Security level will determine the type of encryption/signing to be performed on the
            data as well as the type of keys being exchanged. Must be either ``hmac`` or ``rsa``.
        user_key : str
            String representation of the key from the client. If ``hmac`` this will be the secret
            key used to sign data between client and server; if ``rsa`` this will be the client\'s
            public key.

        Returns
        -------
        response : bottle.HTTPResponse
            Status code and body will determine if the key exchange was successful. If successful,
            the body will contain the single-use token that the client must present along with the
            encoded data. If the security level is ``rsa``, it will also contain the server\'s
            public key.

        Examples
        --------
        >>> app = BottleShip()
        >>> res = app.key_exchange('hmac', '5f04ee43-83bb-46c0-96aa-65a2c5').body
        >>> res
            b'!RWEEi79yWn5BEcse1bmDBBswlcIY2P817ibkZ4UY/kU=?eyJUb2tlbiI6ICI3YTNmMWRlYS0zOGE1LTQ2ODA
            tOTU0Zi1iMjA1OTQwMmVlZGUifQ=='
        >>> bottleship.data_decode(res, '5f04ee43-83bb-46c0-96aa-65a2c5')
            '{"Token": "ce045d43-4e42-4531-9370-8928598b6e26"}'
        '''
        token_json = None
        server_key = None
        user_uuid = str(uuid.uuid4())
        if 'hmac' in security_level and security_level in self.allowed_security:
            server_key = user_key
            token_json = json.dumps({'Token': user_uuid})
            token_json = data_encode(token_json, user_key)