import base64
import pickle
from math import log2
from password_strength import PasswordStats
from getpass import getpass as gp
from secrets import token_urlsafe
from random import randint as rint, SystemRandom as sr
from atexit import register
from gc import collect
from os import urandom, path, remove
from cryptography.fernet import Fernet, InvalidSignature, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

def encrypt_db(password):
    binpass = password.encode()
    salt = urandom(16)
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA512(),
        length=32,
        salt=salt,
        iterations=1500000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(binpass))
    fernet = Fernet(key)

    with open("database/db.dat", "rb") as f:
        data = f.read()

    encr = fernet.encrypt(data)

    with open("database/db.asc", "wb") as f:
        f.write(encr)

    with open("database/db.s", "wb") as f:
        f.write(salt)

def decrypt_db(password):
    with open("database/db.s", "rb") as f:
        salt = f.read()

    binpass = password.encode()
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA512(),
        length=32,
        salt=salt,
        iterations=1500000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(binpass))

    fernet = Fernet(key)

    with open("database/db.asc", "rb") as f:
        encrypted_data = f.read()

    decr = fernet.decrypt(encrypted_data)

    with open("database/db.dat", "wb") as f:
        f.write(decr)
    
    with open("database/db.dat", "rb") as f:
        while True:
            try:
                db = pickle.load(f)
            except EOFError:
                print("\nDatabase loaded.")
                break
    shred()

def shred():
    with open("database/db.dat", "wb") as f:
        for _ in range(5):
            f.seek(0)
            f.write(urandom(path.getsize("database/db.dat")))
        remove("database/db.dat")

def clearmem():
    db = rint(100000000000000000000000000000000000000000000000000000000000, 999999999999999999999999999999999999999999999999999999999999)
    db = None

def gen():
    while True:
        length = int(input("Enter password length (above 8 only): "))
        if length <= 7:
            print("The password is too short. Please enter it again.")
        else:
            break
    pool = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789`~@#$%^&*()-_=+]}[{\"';:.>,</?"
    strength = ''.join(sr().choice(pool) for i in range(length))
    print(strength)

def strength():
    password = gp(prompt = "\nEnter the password to check its strength (hidden for privacy!): ")
    if password in common_passwords:
        print("Your password is in the list of 1000 of the most common passwords. Change your password immediately.")
    else:
        poolsize = 0
        for l in password:
            if l.islower():
                poolsize += 26
            elif l.isupper():
                poolsize += 26
            elif l.isdigit():
                poolsize += 10
            else:
                poolsize += 32
        entropy = log2(poolsize**len(password))
        passstrength = PasswordStats(strength).strength()
        if passstrength < 0.33:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is extremely weak. Please change it.\n")
        elif passstrength >= 0.33 and passstrength < 0.66:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is of fair strength. Please consider changing it to a stronger one.\n")
        elif passstrength >= 0.66 and passstrength < 0.80:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is strong.\n")
        elif passstrength >= 0.80 and passstrength < 0.9:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is incredibly strong.\n")
        else:
            print(f"[PASSWORD STRENGTH]: {passstrength}\n[PASSWORD ENTROPY]: {entropy} bits\nYour password is practically uncrackable.\n")


global db
register(clearmem)

with open("files/strength/common-passwords.txt", "r") as f:
    common_passwords = f.read()

print("\nkeyvault initialized.")

for _ in range(4):
    try:
        if _ != 3:
            password = gp(prompt = "\nEnter your password (hidden for privacy!): ")
            decrypt_db(password)
            break
        else:
            print(" You have exceeded the maximum number of tries.")
            exit()
    except (InvalidSignature, InvalidToken):
        print("Incorrect password.", end = '')

print("keyvault is ready to use! Type 'help' for a list of commands.\n")

while True:
    command = input("> ").lower()
    if command == 'help':
        print("\nUsage:\n\nls - list entires\nrm - remove entry\ngen - generate a password or username\nstrength - check password strength\nedit - edit an entry\nshow/view - view an entry\nfind/search - open search wizard\nhelp - display this message\nversion - print current version\n")
    elif command == 'version':
        print("keyvault: v1.0.0")
    elif command == 'gen':
        gen()
    elif command == 'strength':
        strength()