#!/usr/bin/env python
import argparse
import getpass
import gnupg
import json
import os
import random
import shutil
import string
import site
import sys
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument("--copy-master", nargs=2, help="Export the master.asc to a new file using a different gpg key.")
parser.add_argument("--field", choices=["password", "username"],
                    help="Select which the key field with --get-profile or --modify-profile.")
parser.add_argument("--get-profile", help="Retrieve a field from a profile and store in clipboard. Defaults to password.")
parser.add_argument("--install-completions", action="store_true", help="Install any available shell completions.")
parser.add_argument("--modify-profile", help="Modify a field in a profile. Defaults to username.")
parser.add_argument("--new-master", action="store_true", help="Create a new, encrypted master file.")
parser.add_argument("--new-profile", help="Add a new account profile to the encrypted master file.")
parser.add_argument("--password-length", default=16, type=int, 
                    help="Generate password of length x when creating or modifying a profile.")
parser.add_argument("--remove-profile", help="Remove a stored profile from the encrypted master file.")
args = parser.parse_args()

def copy_master(gpg, key, newkey, path, newpath):
    decrypt = decrypt_master(gpg, key, path)
    master = json.loads(str(decrypt))
    newgpg = gnupg.GPG()
    newgpg.import_keys(key)
    new_encrypt = gpg.encrypt(json.dumps(master).encode("utf-8"), newkey, output=str(newpath))
    if not new_encrypt.ok:
        sys.exit("Encryption failed! Please check that your newkey value is valid")

def create_new_encrypted_master(gpg, key, path):
    if path.is_file():
        print("The master.asc file already exists! This will delete all your data and create a new, blank file!")
        confirm = input("Are you absolutely sure you want to overwrite it? Type 'yes' to proceed.\n")
        if confirm != "yes":
            sys.exit("Quitting.")
    data_dir = path.parent
    data_dir.mkdir(parents=True, exist_ok=True)
    path.touch(exist_ok=True)
    path.write_text("{}")
    stream = open(str(path), "rb")
    encrypt = gpg.encrypt_file(stream, key, output=str(path))
    stream.close()
    if not encrypt.ok:
        sys.exit("Encryption failed! Please check that you inputted a valid email in the config.")

def create_new_profile(gpg, key, profile, length, path):
    decrypt = decrypt_master(gpg, key, path)
    master = json.loads(str(decrypt))
    if profile in master:
        print("The "+profile+" profile already exists in the master.asc file!")
        confirm = input("Are you absolutely sure you want to change its values? Type 'yes' to proceed.\n")
        if confirm != "yes":
            sys.exit("Quitting.")
    username = input("Input username for this profile.\n")
    if length == 0:
        password = getpass.getpass(prompt="Input the password for this profile.\n")
    else:
        password = generate_password(length)
    master[profile] = {"username" : username, "password": password}
    encrypt = gpg.encrypt(json.dumps(master).encode("utf-8"), key, output=str(path))
    if not encrypt.ok:
        sys.exit("Encryption failed! Please check that you inputted a valid email in the config.")

def decrypt_master(gpg, key, path):
    stream = open(str(path), "rb")
    decrypt = gpg.decrypt_file(stream)
    if not decrypt.ok:
        sys.exit("Decryption failed! Please check that you inputted a valid email in the config.")
    return decrypt

def generate_password(length):
    lowercase = string.ascii_lowercase
    uppercase = string.ascii_uppercase
    digits = string.digits
    special = string.punctuation
    special = special.replace("\'","")
    combined = lowercase + uppercase + digits + special

    rand_lower = random.choice(lowercase)
    rand_upper = random.choice(uppercase)
    rand_digit = random.choice(digits)
    rand_special = random.choice(special)

    password = rand_lower + rand_upper + rand_digit + rand_special
    for i in range(length - 4):
        password = password + random.choice(combined)
        password = "".join(random.sample(password, len(password)))
    return password

def get_clipboard_program(path):
    clipboard = read_config_var(path, "clipboard-copy=")
    if clipboard == "":
        sys.exit("Clipboard program not found in config file! Must be a line formatted as 'clipboard-copy=prog --args'!")
    return clipboard

def get_gpg_key(path):
    key = read_config_var(path, "gpg=")
    if key == "":
        sys.exit("Valid gpg email address not found in config file! Must a line formatted as 'gpg=user@example.com'!")
    return key

def get_profile(gpg, key, profile, field, path, clipboard):
    decrypt = decrypt_master(gpg, key, path)
    master = json.loads(str(decrypt))
    if not profile in master:
        sys.exit("The '"+profile+"' profile doesn't exist!")
    value = ""
    if field == "password":
        value = master[profile]["password"]
    elif field == "username":
        value = master[profile]["username"]
    os.system(clipboard+" '"+value+"'")

def install_completions():
   completions_dir = Path(site.getsitepackages()[0]) / "password-sanity" / "completions"
   bash_completion = completions_dir / "bash" / "pws"
   zsh_completion = completions_dir / "zsh" / "_pws"
   bash_dir = Path(sys.prefix) / "share" / "bash-completion" / "completions"
   zsh_dir = Path(sys.prefix) / "share" / "zsh" / "site-functions"
   if bash_dir.is_dir():
      shutil.copy(str(bash_completion), str(bash_dir))
   if zsh_dir.is_dir():
      shutil.copy(str(zsh_completion), str(zsh_dir))

def modify_profile(gpg, key, profile, field, length, path):
    decrypt = decrypt_master(gpg, key, path)
    master = json.loads(str(decrypt))
    if not profile in master:
        sys.exit("The '"+profile+"' profile doesn't exist!")
    print("This will modify the '"+field+"' field in the '"+profile+"' profile.")
    confirm = input("Type 'yes' to proceed.\n")
    if confirm == "yes":
        if field == "username":
            username = input("Type in the new username.\n")
            master[profile][field] = username
        else:
            if length == 0:
                password = getpass.getpass(prompt="Input the password for this profile.\n")
            else:
                password = generate_password(length)
            master[profile][field] = password
    encrypt = gpg.encrypt(json.dumps(master).encode("utf-8"), key, output=str(path))
    if not encrypt.ok:
        sys.exit("Encryption failed! Please check that you inputted a valid email in the config.")

def quit_message():
    config_dir = config_path.parent
    config_dir.mkdir(parents=True, exist_ok=True)
    config_path.touch(exist_ok=True)

def read_config_var(path, var):
    with open(str(path), "r") as f:
        for line in f:
            if line.find(var) != -1:
                value = line.split(var)[1].rstrip()
                return value
    return ""

def remove_profile(gpg, key, profile, path):
    decrypt = decrypt_master(gpg, key, path)
    master = json.loads(str(decrypt))
    if profile in master:
        confirm = input("This will delete the '"+profile+"' profile from the encrypted master. Type 'yes' to proceed.\n")
        if confirm == "yes":
            del master[profile]
            encrypt = gpg.encrypt(json.dumps(master).encode("utf-8"), key, output=str(path))
            if not encrypt.ok:
                sys.exit("Encryption failed! Please check that you inputted a valid email in the config.")
        else:
            sys.exit("Quitting.")
    else:
        sys.exit("The '"+profile+"' profile was not found.")

def main(args):
    config_path = Path.home() / ".config" / "pws" / "config"
    data_path = Path.home() / ".local" / "share" / "pws" / "master.asc"

    if args.install_completions:
        install_completions()
        sys.exit()

    if not config_path.is_file():
        sys.exit("Config file in ~/.config/pws/config does not yet exist. Please create it with the right variables.")

    key = get_gpg_key(config_path)
    clipboard = get_clipboard_program(config_path)
    gpg = gnupg.GPG()
    gpg.import_keys(key)

    if not data_path.is_file():
        confirm = input("Encrypted master file does not exist. Would you like to create it now? Type 'yes' to proceed.\n")
        if confirm == 'yes':
            create_new_encrypted_master(gpg, key, data_path)
            sys.exit()

    if args.copy_master:
        newkey = args.copy_master[0]
        newpath = args.copy_master[1]
        copy_master(gpg, key, newkey, data_path, newpath)
    elif args.get_profile:
        if args.field == None:
            args.field = "password"
        get_profile(gpg, key, args.get_profile, args.field, data_path, clipboard)
    elif args.modify_profile:
        if args.field == None:
            args.field = "username"
        modify_profile(gpg, key, args.modify_profile, args.field, args.password_length, data_path)
    elif args.new_master:
        create_new_encrypted_master(gpg, key, data_path)
    elif args.new_profile:
        create_new_profile(gpg, key, args.new_profile, args.password_length, data_path)
    elif args.remove_profile:
        remove_profile(gpg, key, args.remove_profile, data_path)
    else:
        parser.print_help()

if __name__ == "__main__":
    main(args)
