#!/usr/bin/env python
import argparse
import getpass
import gnupg
import json
import os
import pyperclip
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", "profile-name"],
                    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("--list-profiles", action="store_true", help="List all saved profiles.")
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 main(args):
    config_path = get_config_path()
    data_path = get_data_path()

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

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

    key = get_gpg_key(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 = Path(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)
    elif args.list_profiles:
        list_profiles(gpg, key, data_path)
    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()

def copy_master(gpg, key, newkey, path, newpath):
    master = decrypt_master(gpg, key, path)
    gpg.import_keys(newkey)
    encrypt_master(gpg, newkey, newpath, master)

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 your gpg key is correct, and you have write permissions.")

def create_new_profile(gpg, key, profile, length, path):
    master = decrypt_master(gpg, key, path)
    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_master(gpg, key, path, master)

def decrypt_master(gpg, key, path):
    stream = open(str(path), "rb")
    decrypt = gpg.decrypt_file(stream)
    stream.close()
    if not decrypt.ok:
        sys.exit("Decryption failed! Please check that your gpg key is correct.")
    master = json.loads(str(decrypt))
    return master

def encrypt_master(gpg, key, path, master):
    tmp_path = path.parent / "master_tmp.asc"
    encrypt = gpg.encrypt(json.dumps(master).encode("utf-8"), key, output=str(tmp_path))
    if not encrypt.ok:
        sys.exit("Encryption failed! Please check that your gpg key is correct, and you have write permissions.")
    if path.is_file():
        os.remove(str(path))
    os.rename(str(tmp_path), str(path))

def generate_password(length):
    lowercase = string.ascii_lowercase
    uppercase = string.ascii_uppercase
    digits = string.digits
    special = string.punctuation
    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_config_path():
    if sys.platform == "win32":
        return Path(os.environ.get("APPDATA")) /"pws" / "config"
    else:
        return Path.home() / ".config" / "pws" / "config"

def get_data_path():
    if sys.platform == "win32":
        return Path(os.environ.get("APPDATA")) / "pws" / "master.asc"
    else:
        return Path.home() / ".local" / "share" / "pws" / "master.asc"

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):
    master = decrypt_master(gpg, key, path)
    if not profile in master:
        sys.exit("The '"+profile+"' profile doesn't exist!")
    value = ""
    if field == "password":
        value = master[profile]["password"]
    elif field == "profile-name":
        value = profile
    elif field == "username":
        value = master[profile]["username"]
    pyperclip.copy(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 list_profiles(gpg, key, path):
    master = decrypt_master(gpg, key, path)
    for profile in sorted(master.keys()):
        print(profile)

def modify_profile(gpg, key, profile, field, length, path):
    master = decrypt_master(gpg, key, path)
    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 == "password":
            if length == 0:
                password = getpass.getpass(prompt="Input the password for this profile.\n")
            else:
                password = generate_password(length)
            master[profile][field] = password
        elif field == "profile-name":
            profile_name = input("Type in the new profile-name.\n")
            master[profile_name] = master[profile]
            del master[profile]
        elif field == "username":
            username = input("Type in the new username.\n")
            master[profile][field] = username
    encrypt_master(gpg, key, path, master)

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):
    master = decrypt_master(gpg, key, path)
    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_master(gpg, key, path, master)
        else:
            sys.exit("Quitting.")
    else:
        sys.exit("The '"+profile+"' profile was not found.")

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