import base64
import json
import uuid
import hashlib
from functools import wraps
from flask import request, make_response, jsonify, session
from Crypto import Random
from Crypto.Cipher import AES
# from Crypto import Random
# from Crypto.Cipher import AES
from datetime import datetime

from scripts.config import app_configuration as config, app_constants, app_configuration

max_age_in_mins = config.cookie_max_age


class AESCipher(object):
    """
    A classical AES Cipher. Can use any size of data and any size of password thanks to padding.
    Also ensure the coherence and the type of the data with a unicode to byte converter.
    """

    def __init__(self, key):
        self.bs = 16
        self.key = AESCipher.str_to_bytes(key)

    @staticmethod
    def str_to_bytes(data):
        u_type = type(b''.decode('utf8'))
        if isinstance(data, u_type):
            return data.encode('utf8')
        return data

    def _pad(self, s):
        return s + (self.bs - len(s) % self.bs) * AESCipher.str_to_bytes(chr(self.bs - len(s) % self.bs))

    @staticmethod
    def _unpad(s):
        return s[:-ord(s[len(s) - 1:])]

    def encrypt(self, raw):
        raw = self._pad(AESCipher.str_to_bytes(raw))
        iv = Random.new().read(AES.block_size)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return base64.b64encode(iv + cipher.encrypt(raw)).decode('utf-8')

    def decrypt(self, enc):
        enc = base64.b64decode(enc)
        iv = enc[:AES.block_size]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        data = self._unpad(cipher.decrypt(enc[AES.block_size:]))
        return data.decode('utf-8')


def create_cookie(request_json, request, user_id=None):
    """
    This method is to create a cookie
    """
    status = False
    cookie_encoded_str = ""
    # user_id = ""
    try:
        if user_id is None:
            user_id = request.cookies.get("user_id")
        client_ip = request.headers.get('X-Forwarded-For').split(',')[0].split(':')[0]
        cookie_string = (
                user_id + '^' +
                client_ip + "^" +
                request.headers.get('User-Agent') + "^" + datetime.now().strftime('%Y-%m-%dT%H:%M:%SZ'))
        cookie_encoded_str = AESCipher(app_constants.KEYS.cookie_encryption_private_key).encrypt(cookie_string)
        return True, cookie_encoded_str
    except Exception as e:
        # traceback.print_exc()
        return status, cookie_encoded_str


def validate_cookie(cookie_enc_str, request):
    """
    This method is to validate the cookie
    """
    status = False
    encrypted_str = None
    ip_check = True
    user_id = request.cookies['user_id']
    try:
        headers = request.headers
        decrypted_string = AESCipher(app_constants.KEYS.cookie_encryption_private_key).decrypt(cookie_enc_str)
        decrypted_data = decrypted_string.split('^')
        client_ip = headers.get('X-Forwarded-For').split(',')[0].split(':')[0]
        user_agent = request.headers.get('User-Agent')
        time_diff = (datetime.now() - datetime.strptime(decrypted_data[3], "%Y-%m-%dT%H:%M:%SZ")).total_seconds()
        if time_diff > max_age_in_mins * 60:
            raise Exception("Session Expired")
        #Commented the existing code for Hotfix - Session Logout happening due to Validation of IP happening in Client Networks Fixed by VigneshRavishankar
        #if decrypted_data[0] == user_id and \
        #         decrypted_data[1] == client_ip and \
        #         decrypted_data[2] == request.headers.get('User-Agent'):
        #     encrypted_str = decrypted_data
        #     return True, encrypted_str
        # print(f'Input User: {decrypted_data[0]} | Validate User: {user_id}')
        # print(f'Input IP: {decrypted_data[1]} | Validate IP: {client_ip}')
        # print(f'Input Agent: {decrypted_data[2]} | Validate Agent: {user_agent}')
        
        if decrypted_data[0] == user_id:
            user_check = True
        else:
            raise Exception('Invalid Cookie (User validation failed)')

        if app_configuration.ip_check:
            if decrypted_data[1] == client_ip:
                ip_check = True
            else:
                raise Exception('Invalid Cookie (Agent validation failed)')

        if decrypted_data[2] == user_agent:
            agent_check = True
        else:
            raise Exception('Invalid Cookie (Agent validation failed)')
        if user_check and ip_check and agent_check:
            encrypted_str = decrypted_data
            return True, encrypted_str
        else:
            raise Exception("Invalid cookie")
    except Exception as e:
        print("Error : {}".format(str(e)))
        return status, encrypted_str


def apply_encryption(func):
    """
    This function is to be use as a decorator only,
    This function works like a wrapper which will decrypt the data before sending to the function and also encrypt
    the data before returning.
    :param func:
    :return: wrapper
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        if config.SERVICE_ENABLE_SECURITY is True:
            session_id = str(request.cookies['session_id'])[:16]
            data = request.data
            data = data.decode()
            if 'session_id' in request.cookies:
                status, encoded_str = validate_cookie(request.cookies['session_id'], request)
            else:
                return make_response(json.dumps({"status": "session_invalid", "message": "Invalid Session"}), 401)
            if status:
                request.data = json.loads(AESCipher(session_id).decrypt(data))
                response = func(*args, **kwargs)
                data = AESCipher(session_id).encrypt(json.dumps(response))
                return data
            else:
                return make_response(json.dumps({"status": "session_invalid", "message": "Invalid Session"}), 401)

        else:
            # if "session_id" not in session:
            #     return make_response(json.dumps({"status": "session_invalid", "message": "Invalid Session"}), 401)

            request.data = json.loads(request.data.decode())

            if 'session_id' in request.cookies:
                status, encoded_str = validate_cookie(request.cookies['session_id'], request)
            else:
                return make_response(json.dumps({"status": "session_invalid", "message": "Invalid Session"}), 401)
            if status:
                response = func(*args, **kwargs)
                if str(request.url_rule).endswith('reset_password'):
                    return response
                else:
                    response = make_response(json.dumps(response), 200)
                    response.headers['X-Content-Type-Options'] = 'nosniff'
                    response.headers['X-Frame-Options'] = 'SAMEORIGIN'
                    response.headers['Cache-Control'] = 'no-store'
                    cookie_status, secure_cookie = create_cookie(request.data, request)
                    # response.set_cookie('session_id', request.cookies['session_id'], max_age=max_age_in_mins * 60 ,httponly=app_configuration.http_flag)
                    if cookie_status:
                        response.set_cookie('session_id', secure_cookie, httponly=app_configuration.http_flag)
                        return response
                    else:
                        raise Exception("Failed to set cookie")
            else:
                return make_response(json.dumps({"status": "session_invalid", "message": "Invalid Session"}), 401)

    return wrapper


def login_wrap(func):
    """
    This function is used as a decorator only for login services
    :param func:
    :return: wrapper
    """

    @wraps(func)
    def wrap_login(*args, **kwargs):
        if config.SERVICE_ENABLE_SECURITY is True:
            # session_id_1 = str(uuid.uuid4()).replace("-", "")
            # session["session_id"] = session_id_1
            session_id = str(session["session_id"])
            response = func(*args, **kwargs)
            resp = make_response(response, 200)
            resp.set_cookie("session_id", session_id, max_age=max_age_in_mins * 60)
            return resp
        else:
            # session_id_1 = str(uuid.uuid4()).replace("-", "")
            # session["session_id"] = session_id_1
            session_id = str(session["session_id"])
            request.data = json.loads(request.data.decode())
            response = func(*args, **kwargs)
            response = make_response(response, 200)
            response.set_cookie("session_id", session_id, max_age=max_age_in_mins * 60)
            return response

    return wrap_login


class CipherEncryption(object):

    def __init__(self, key):
        self.bs = AES.block_size
        self.key = hashlib.sha256(key.encode()).digest()

    def encrypt(self, raw):
        raw = self._pad(raw)
        iv = Random.new().read(AES.block_size)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return base64.b64encode(iv + cipher.encrypt(raw.encode()))

    def decrypt(self, enc):
        enc = base64.b64decode(enc)
        iv = enc[:AES.block_size]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return self._unpad(cipher.decrypt(enc[AES.block_size:])).decode('utf-8')

    def _pad(self, s):
        return s + (self.bs - len(s) % self.bs) * chr(self.bs - len(s) % self.bs)

    @staticmethod
    def _unpad(s):
        return s[:-ord(s[len(s) - 1:])]
