#!/usr/bin/env python3

import json
import logging
import os
import shutil
import sys
from base64 import b64encode

from Crypto.Hash import SHA256
from Crypto.Protocol.KDF import bcrypt, bcrypt_check
from Crypto.PublicKey import RSA
from Crypto.PublicKey.RSA import RsaKey

auth_logger = logging.getLogger('AUTH APPLICATION ')
auth_logger.root.setLevel(logging.INFO)

class Authetication:
    ABSOLUTE_PATH = os.path.abspath(os.path.dirname(sys.argv[0]))
    HOME_DIRECTORY_LOCATION = ABSOLUTE_PATH + os.path.sep + "home"
    CONFIG_DIRECTORY_LOCATION = ABSOLUTE_PATH + os.path.sep + "config"
    CONFIG_FILE_LOCATION = ABSOLUTE_PATH + os.path.sep + "config" + os.path.sep + "config.txt"
    PRIVATE_KEY_DIRECTORY_LOCATION = CONFIG_DIRECTORY_LOCATION + os.path.sep + "private_keys"
    USER_INDEX = 0


    def __init__(self):
        if not os.path.isdir(self.HOME_DIRECTORY_LOCATION):
            os.mkdir(self.HOME_DIRECTORY_LOCATION)

        if not os.path.isdir(self.CONFIG_DIRECTORY_LOCATION):
            os.mkdir(self.CONFIG_DIRECTORY_LOCATION)

        if not os.path.isdir(self.PRIVATE_KEY_DIRECTORY_LOCATION):
            os.mkdir(self.PRIVATE_KEY_DIRECTORY_LOCATION)

        if not os.path.isfile(self.CONFIG_FILE_LOCATION) or os.stat(
                self.CONFIG_FILE_LOCATION).st_size == 0:
            data = {'index': 0, 'user': []}
            with open(self.CONFIG_FILE_LOCATION, 'w+') as outfile:
                json.dump(data, outfile)


    def login(self, username: str, password: str) -> str:
        with open(self.CONFIG_FILE_LOCATION) as json_file:
            data = json.load(json_file)

        for user in data['user']:
            if username == user['username']:
                b64pwd = b64encode(SHA256.new(password.encode('utf-8')).digest())
                try:
                    b64pwd = b64encode(SHA256.new(password.encode('utf-8')).digest())
                    bcrypt_check(b64pwd, user['password'].encode('utf-8'))
                    auth_logger.debug("User logged in: " + username)
                    return user['homeDir']
                except ValueError:
                    auth_logger.debug("User NOT logged in: " + username)
                    return ''
        else:
            auth_logger.error("User not found")
            return ''

    def checkUserExists(self, username: str) -> bool:
        with open(self.CONFIG_FILE_LOCATION) as json_file:
            data = json.load(json_file)

        for user in data['user']:
            if username == user['username']:
                return True
        return False


    def initConfig(self):
        data = {'index': 0, 'user': []}
        with open(self.CONFIG_FILE_LOCATION, 'w+') as outfile:
            json.dump(data, outfile)

        shutil.rmtree(self.HOME_DIRECTORY_LOCATION)
        os.mkdir(self.HOME_DIRECTORY_LOCATION)
        shutil.rmtree(self.PRIVATE_KEY_DIRECTORY_LOCATION)
        os.mkdir(self.PRIVATE_KEY_DIRECTORY_LOCATION)


    def saveUser(self, username: str, password: str) -> bool:
        bytePass = password.encode('utf-8')
        b64pwd = b64encode(SHA256.new(bytePass).digest())
        bcrypt_hash = bcrypt(b64pwd, 12)

        with open(self.CONFIG_FILE_LOCATION) as json_file:
            data = json.load(json_file)

        if self.checkUserExists(username):
            auth_logger.debug("User NOT saved! This username already exists!")
            return False
        else:
            if not os.path.isdir(self.HOME_DIRECTORY_LOCATION + os.path.sep + str(data['index'] + 1)):
                data['index'] = data['index'] + 1
                user = {
                    'username': username,
                    'password': bcrypt_hash.decode('utf-8'),
                    'homeDir': str(data['index']),
                    'publicKey': ''
                }

                ##Create user HOME directory with index as name
                os.mkdir(self.HOME_DIRECTORY_LOCATION + os.path.sep + str(data['index']))

                ##Save user data
                data['user'].append(user)
                with open(self.CONFIG_FILE_LOCATION, 'w') as outfile:
                    json.dump(data, outfile)

                auth_logger.debug("User saved!")
            else:
                auth_logger.debug("User NOT saved! Home directory already exists")

            return True


    def loadUserPublicKeys(self) -> dict:
        with open(self.CONFIG_FILE_LOCATION) as json_file:
            data = json.load(json_file)

        dictionary: dict

        for user in data['user']:
            key = data['publicKey']
            key = bytes.fromhex(key)
            try:
                rsaKey = RSA.import_key(key)
                dictionary[user['username']] = rsaKey
            except ValueError:
                print('Invalid server public key!')

        return dictionary


    def loadServerPrivateKey(self, passphrase: str) -> RsaKey:
        with open(self.CONFIG_FILE_LOCATION) as json_file:
            data = json.load(json_file)

        key = data['serverPrivateKey']
        key = bytes.fromhex(key)
        try:
            rsaKey = RSA.import_key(key, passphrase)
        except ValueError:
            print('Invalid server private key!')
        return rsaKey