#!/usr/bin/python3

# external modules

from os import environ, mkdir, remove, system
from shutil import copy, move, rmtree
from sys import argv, exit as s_exit
import random
import sshync
import string


# utility functions

def replace_line(file_name, line_num, text):  # replaces text in a given line with different text
    _lines = open(file_name, 'r').readlines()
    _lines[line_num] = text
    open(file_name, 'w').writelines(_lines)


def shm_gen():  # generates random folder in /dev/shm for security and returns random names for the folder and temp file
    _shm_folder_gen = ''.join(random.SystemRandom().choice(string.ascii_letters + string.digits)
                              for _ in range(random.randint(10, 30)))
    _shm_entry_gen = ''.join(random.SystemRandom().choice(string.ascii_letters + string.digits)
                             for _ in range(random.randint(10, 30)))
    mkdir(f"/dev/shm/{_shm_folder_gen}", 0o700)
    return _shm_folder_gen, _shm_entry_gen


def encrypt(_shm_folder, _shm_entry, _location):
    system(f"gpg -r {str(gpg_id)} -e '/dev/shm/{_shm_folder}/{_shm_entry}'")
    move(f"/dev/shm/{_shm_folder}/{_shm_entry}.gpg", f"{directory}{_location}.gpg")
    rmtree(f"/dev/shm/{_shm_folder}")


def edit_note(_shm_folder, _shm_entry):
    lines = open(f"/dev/shm/{_shm_folder}/{_shm_entry}").readlines()
    open(f"/dev/shm/{_shm_folder}/{_shm_entry}", 'w').writelines(lines[0:3])
    open(f"/dev/shm/{_shm_folder}/{_shm_entry}-n", 'w').writelines(lines[4:-2])
    system(f"nano /dev/shm/{_shm_folder}/{_shm_entry}-n")
    edit_notes = open(f"/dev/shm/{_shm_folder}/{_shm_entry}-n").read()
    open(f"/dev/shm/{_shm_folder}/{_shm_entry}", 'a').write('\n' + edit_notes + '\n\n')


# argument-specific functions


def tweak():  # runs configuration wizard
    # storage directory creation
    try:
        mkdir(f"/home/{environ.get('USER')}/.password-pasture", 0o700)
    except FileExistsError:
        pass

    # device type configuration
    _device_type = input('\nIs this installation for a client or for a server? (C/s) ')
    if _device_type.lower() == 's':
        open("/var/lib/sshyp/sshyp-device", 'w').write(_device_type)
        copy('/usr/bin/sshync.py', f"/home/{environ.get('USER')}/")
        print('\nMake sure the ssh service is running and properly configured.\n\nConfiguration complete!\n')
        s_exit()
    else:
        # device type configuration
        _device_type = 'c'  # ensure device type flag is properly set
        open("/var/lib/sshyp/sshyp-device", 'w').write(_device_type)

        # gpg configuration
        _gpg_id = input('\nsshyp requires the use of a unique gpg key. Do you already have one that you are '
                        'willing to use? (y/N)')
        if _gpg_id.lower() != 'y':
            print('\nA unique gpg key has been generated for you.')
            system('gpg --full-generate-key')
        _gpg_id = str(input('\nPlease input the ID of your gpg key: '))
        open("/var/lib/sshyp/sshyp-gpg", 'w').write(_gpg_id + '\n')

        # lock file generation
        try:
            open('/var/lib/sshyp/lock', 'x').write('')
        except FileExistsError:
            pass
        system(f"gpg -r {str(_gpg_id)} -e /var/lib/sshyp/lock")
        remove('/var/lib/sshyp/lock')

        # ssh key configuration
        _ssh_gen = (input('\nMake sure the ssh service on the remote server is running and properly configured.'
                          '\n\nSync support requires a unique ssh key. Would you like to have this automatically '
                          'generated? (Y/n) '))
        if _ssh_gen.lower() != 'n':
            system('ssh-keygen -t ed25519 -f ~/.ssh/sshyp')
        else:
            print(f"\nEnsure that the key file you are using is located at /home/{environ.get('USER')}/.ssh/sshyp")

        # ssh ip+port configuration
        print('\nExample input: 10.10.10.10:9999\n')
        _ip_port = str(input('What is the IP and ssh port of the remote server? '))
        _ip, _sep, _port = _ip_port.partition(':')

        # ssh user configuration
        _username_ssh = str(input('\nWhat is the username of the remote server? '))
        print(f"\nEntry directory: /home/{_username_ssh}/.password-pasture/")

        # sshync profile generation
        sshync.make_profile('/var/lib/sshyp/sshyp.sshync', f"/home/{environ.get('USER')}/.password-pasture/",
                            f"/home/{_username_ssh}/.password-pasture/", f"/home/{environ.get('USER')}/.ssh/sshyp",
                            _ip, _port, _username_ssh)
        print('\nConfiguration complete!\n')


def helper():  # displays help menu
    print('\nsshyp  Copyright (C) 2021  Randall Winkhart')
    print("This is free software, and you are welcome to redistribute it under certain conditions;\nthis program comes "
          "with ABSOLUTELY NO WARRANTY;\ntype `sshyp license' for details.")
    print('\nUSAGE: sshyp [/<entry name>] [OPTION] [FLAG]\n')
    print('Options: ')
    print('help/--help/-h           bring up this menu')
    print('version/-v               display sshyp version info')
    print('tweak                    configure sshyp')
    print('add                      add an entry')
    print('gen                      generate a new password')
    print('edit                     edit an existing entry')
    print('copy                     copy details of an entry to your clipboard')
    print('shear/-rm                delete an existing entry')
    print('sync/-s                  manually sync the entry directory via sshync\n')
    print('Flags:')
    print('add:')
    print(' password/-p             add a password entry')
    print(' note/-n                 add a note entry')
    print(' folder/-f               add a new folder for entries')
    print('edit:')
    print(' rename/relocate/-r      rename or relocate an entry')
    print(' username/-u             change the username of an entry')
    print(' password/-p             change the password of an entry')
    print(' note/-n                 change the note attached to an entry')
    print(' url/-l                  change the url attached to an entry')
    print('copy:')
    print(' username/-u             copy the username of an entry to your clipboard')
    print(' password/-p             copy the password of an entry to your clipboard')
    print(' url/-l                  copy the URL of an entry to your clipboard')
    print(' note/-n                 copy the note of an entry to your clipboard')
    print('gen:')
    print(' update/-u               generate a password for an existing entry\n')
    print("Tip: You can quickly read an entry with 'sshyp /<entry name>'!")
    print("When specifying entry names, do not include the file extension! '.gpg' is assumed!\n")


def version():  # displays version information
    print('\n////////////////////////////////////////////////////////')
    print('/                                                      /')
    print('/     sshyp  Copyright (C) 2021  Randall Winkhart      /')
    print('/                                                      /')
    print('/              Version 2022.01.10.fr2.3                /')
    print('/                 The Lighter Update                   /')
    print('/                                                      /')
    print('////////////////////////////////////////////////////////\n')


def no_arg():  # displays a list of entries and gives an option to select one for viewing
    print("\nFor a list of usable commands, run 'sshyp help'.\n")
    system(f"cd {directory}; ls -1 *")
    _read_entry = str(input('\nEntry to read: '))
    system(f"gpg -d '{directory}{_read_entry.replace('/', '', 1)}.gpg'")
    print()


def show_license():  # displays licensing information
    print('\nThis program is free software: you can redistribute it and/or modify it under the terms of the GNU General'
          '\nPublic License as published by the Free Software Foundation, either version 3 of the License,\nor '
          '(at your option) any later version.\n'
          '\nThis program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;\nwithout even '
          'the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.\nSee the GNU General Public '
          'License for more details.\n\nhttps://opensource.org/licenses/GPL-3.0\n')


def read_shortcut():  # shortcut to quickly read an entry
    system(f"gpg -d '{directory}{argument.replace('/', '', 1)}.gpg'")


def sync():  # calls sshync to sync changes to the user's server
    print('\nSyncing entries with the server device...\n')
    # set permissions before uploading
    system('find ' + directory + ' -type d -exec chmod -R 700 {} +')
    system('find ' + directory + ' -type f -exec chmod -R 600 {} +')
    sshync.run_profile('/var/lib/sshyp/sshyp.sshync')


def add_info():  # displays 'add' syntax
    print('\nUSAGE: sshyp add [FLAG]')
    print('Flags:')
    print('add:')
    print(' password/-p             add a password entry')
    print(' note/-n                 add a note entry')
    print(' folder/-f               add a new folder for entries\n')


def add_entry():  # adds a new entry
    _entry_name = str(input('\nName of entry: '))
    if _entry_name.startswith('/'):
        _entry_name = _entry_name.replace('/', '', 1)
    _shm_folder, _shm_entry = shm_gen()
    if argument == 'add note' or argument == 'add -n':
        system(f"nano /dev/shm/{_shm_folder}/{_shm_entry}-n")
        _notes = open(f"/dev/shm/{_shm_folder}/{_shm_entry}-n", 'r').read()
        open(f"/dev/shm/{_shm_folder}/{_shm_entry}", 'w').writelines('\n\n\n\n' + _notes + '\n\n')
    if argument == 'add password' or argument == 'add -p':
        _username = str(input('Username: '))
        _password = str(input('Password: '))
        _url = str(input('URL: '))
        _add_note = input('Add a note to this entry? (y/N) ')
        if _add_note.lower() == 'y':
            system(f"nano /dev/shm/{_shm_folder}/{_shm_entry}-n")
            _notes = open(f"/dev/shm/{_shm_folder}/{_shm_entry}-n", 'r').read()
        else:
            _notes = ''
        open(f"/dev/shm/{_shm_folder}/{_shm_entry}", 'w').writelines(_username + '\n' + _password + '\n' + _url + '\n\n'
                                                                     + _notes + '\n\n')
    encrypt(_shm_folder, _shm_entry, _entry_name)
    system(f"gpg -d '{directory}{_entry_name}.gpg'")


def add_folder():  # creates a new folder
    _folder_name = str(input('Name of new folder: '))
    mkdir(directory + _folder_name, 0o700)
    system(f"ssh -p {port} {username_ssh}@{ip} \"mkdir -p '{directory_ssh}{_folder_name}'\"")


def edit_info():  # displays 'edit' syntax
    print('\nUSAGE: sshyp edit [FLAG]')
    print('Flags:')
    print('edit:')
    print(' rename/relocate/-r      rename or relocate an entry')
    print(' username/-u             change the username of an entry')
    print(' password/-p             change the password of an entry')
    print(' url/-l                  change the url attached to an entry')
    print(' note/-n                 change the note attached to an entry\n')


def rename():  # renames an entry or folder TODO delete original from server to prevent re-downloading
    print()
    system(f"cd {directory}; ls -1 *")  # TODO replace with new list
    _entry_name = str(input('\nWhat would you like to rename/relocate? '))
    _new_name = str(input('New Name: '))
    if _entry_name.startswith('/'):
        if _entry_name.replace('/', '', 1).__contains__('/'):
            _folder, _sep, _folder_entry = _entry_name.replace('/', '', 1).partition('/')
            move(f"{directory}{_folder}/{_folder_entry}.gpg", f"{directory}/{_new_name}.gpg")
        else:
            move(f"{directory}{_entry_name.replace('/', '', 1)}.gpg", f"{directory}{_new_name.replace('/', '', 1)}.gpg")
    else:
        move(f"{directory}{_entry_name}.gpg", f"{directory}{_new_name}.gpg")


def edit():  # edits the contents of an entry
    print()
    system(f"cd {directory}; ls -1 *")  # TODO replace with new list
    _entry_name = str(input('\nWhat file would you like to edit? '))
    if _entry_name.startswith('/'):
        _entry_name = _entry_name.replace('/', '', 1)
    _shm_folder, _shm_entry = shm_gen()
    system(f"gpg -d --output /dev/shm/{_shm_folder}/{_shm_entry} '{directory}{_entry_name}.gpg'")
    if argument == 'edit username' or argument == 'edit -u':
        _detail = str(input('New Username: '))
        replace_line(f"/dev/shm/{_shm_folder}/{_shm_entry}", 0, _detail + '\n')
    if argument == 'edit password' or argument == 'edit -p':
        _detail = str(input('New Password: '))
        replace_line(f"/dev/shm/{_shm_folder}/{_shm_entry}", 1, _detail + '\n')
    if argument == 'edit url' or argument == 'edit -l':
        _detail = str(input('New URL: '))
        replace_line(f"/dev/shm/{_shm_folder}/{_shm_entry}", 2, _detail + '\n')
    if argument == 'edit note' or argument == 'edit -n':
        edit_note(_shm_folder, _shm_entry)
    remove(f"{directory}{_entry_name}.gpg")
    encrypt(_shm_folder, _shm_entry, _entry_name)
    system(f"gpg -d '{directory}{_entry_name}.gpg'")


def gen():  # generates a password for a new or an existing entry
    _username, _url, _notes = None, None, None  # sets base-line values to avoid errors
    _shm_folder, _shm_entry = shm_gen()
    if argument == 'gen update' or argument == 'gen -u':
        print()
        system(f"cd {directory}; ls -1 *")  # TODO replace with new list
    _entry_name = str(input('\nName of service: '))
    if _entry_name.startswith('/'):
        _entry_name = _entry_name.replace('/', '', 1)
    if argument != 'gen update' and argument != 'gen -u':
        _username = str(input('Username: '))
    _pass_length = int(input('How many characters should be in the password? '))
    _pass_type = str(input('Should the password be simple (for compatibility) or complex (for security)? (s/C) '))
    if _pass_type != 's':
        _pass_type = string.ascii_letters + string.digits + string.punctuation
    if _pass_type == 's':
        _pass_type = string.ascii_letters + string.digits
    _pass_gen = ''.join(random.SystemRandom().choice(_pass_type) for _ in range(_pass_length))
    if argument != 'gen update' and argument != 'gen -u':
        _url = str(input('URL: '))
        _notes = str(input('Additional notes: '))  # TODO update w/new notes system
    if argument != 'gen update' and argument != 'gen -u':
        open(f"/dev/shm/{_shm_folder}/{_shm_entry}", 'w').writelines(_username + '\n' + _pass_gen + '\n' + _url +
                                                                     '\n\n' + _notes + '\n\n')
    if argument != 'gen update' and argument != 'gen -u':
        encrypt(_shm_folder, _shm_entry, _entry_name)
        system(f"gpg -d '{directory}{_entry_name}.gpg'")
    else:
        system(f"gpg -d --output /dev/shm/{_shm_folder}/{_shm_entry} '{directory}{_entry_name}.gpg'")
        replace_line(f"/dev/shm/{_shm_folder}/{_shm_entry}", 1, _pass_gen + '\n')
        remove(f"{directory}{_entry_name}.gpg")
        encrypt(_shm_folder, _shm_entry, _entry_name)
        system(f"gpg -d '{directory}{_entry_name}.gpg'")


def copy_info():  # displays 'copy' syntax
    print('\nUSAGE: sshyp copy [FLAG]')
    print('Flags:')
    print('copy:')
    print(' username/-u             copy the username of an entry to your clipboard')
    print(' password/-p             copy the password of an entry to your clipboard')
    print(' url/-l                  copy the URL of an entry to your clipboard')
    print(' note/-n                 copy the note of an entry to your clipboard\n')


def copy_data():  # copies a specified field of an entry to the clipboard
    print()
    system(f"cd {directory}; ls -1 *")  # TODO replace with new list
    _read_entry = str(input('\nEntry to copy: '))
    _shm_folder, _shm_entry = shm_gen()
    system(f"gpg -d --output /dev/shm/{_shm_folder}/{_shm_entry} '{directory}{_read_entry}.gpg'")
    _copy_line = open(f"/dev/shm/{_shm_folder}/{_shm_entry}", 'r').readlines()
    if environ.get('WAYLAND_DISPLAY') == 'wayland-0':  # checks if the user is using Wayland or X for clipboard
        if argument == 'copy username' or argument == 'copy -u':
            system('wl-copy ' + "'" + _copy_line[0].replace('\n', '').replace("'", "'\\''") + "'")
        elif argument == 'copy password' or argument == 'copy -p':
            system('wl-copy ' + "'" + _copy_line[1].replace('\n', '').replace("'", "'\\''") + "'")
        elif argument == 'copy url' or argument == 'copy -l':
            system('wl-copy ' + "'" + _copy_line[2].replace('\n', '').replace("'", "'\\''") + "'")
        elif argument == 'copy note' or argument == 'copy -n':
            system('wl-copy ' + "'" + _copy_line[4].replace('\n', '').replace("'", "'\\''") + "'")
    else:
        if argument == 'copy username' or argument == 'copy -u':
            system('echo -n ' + "'" + _copy_line[0].replace('\n', '').replace("'", "'\\''") + "'" + ' | xclip -sel c')
        elif argument == 'copy password' or argument == 'copy -p':
            system('echo -n ' + "'" + _copy_line[1].replace('\n', '').replace("'", "'\\''") + "'" + ' | xclip -sel c')
        elif argument == 'copy url' or argument == 'copy -l':
            system('echo -n ' + "'" + _copy_line[2].replace('\n', '').replace("'", "'\\''") + "'" + ' | xclip -sel c')
        elif argument == 'copy note' or argument == 'copy -n':
            system('echo -n ' + "'" + _copy_line[4].replace('\n', '').replace("'", "'\\''") + "'" + ' | xclip -sel c')
    rmtree(f"/dev/shm/{_shm_folder}")


def remove_data():  # deletes an entry or folder from both the client and the server
    print()
    system(f"cd {directory}; ls -1 *")  # TODO replace with new list
    _entry_name = str(input('\nWhat would you like to delete? '))
    if _entry_name.startswith('/'):
        _entry_name = _entry_name.replace('/', '', 1)
    try:
        system('gpg -d --output /dev/shm/sshyp.lock /var/lib/sshyp/lock.gpg')
        _unlock = open('/dev/shm/sshyp.lock', 'r').readlines()
        remove('/dev/shm/sshyp.lock')
    except FileNotFoundError:
        print('\nAccess denied.\n')
        s_exit()
    if device_type == 'c':
        system(f"ssh -p {port} {username_ssh}@{ip} \"rm -rf '{directory_ssh}{_entry_name}.gpg'\"")
    remove(f"{directory}{_entry_name}.gpg")


# program start sequence

# retrieve typed argument
argument_list = argv
i, argument = 0, ''
for _ in argument_list:
    while i < len(argument_list) - 1:
        i += 1
        argument += argument_list[i] + ' '
argument = argument[:-1]

# import saved userdata
if argument != 'tweak':
    device_type = ''
    try:
        device_type = open('/var/lib/sshyp/sshyp-device').read().strip()
        if device_type == 'c':
            gpg_id = open('/var/lib/sshyp/sshyp-gpg').read().strip()
            ssh_info = sshync.get_profile('/var/lib/sshyp/sshyp.sshync')
            username_ssh = ssh_info[0].replace('\n', '')
            ip = ssh_info[1].replace('\n', '')
            port = ssh_info[2].replace('\n', '')
            directory = str(ssh_info[3].replace('\n', ''))
            directory_ssh = '/home/' + username_ssh + '/.password-pasture/'
    except FileNotFoundError:
        if device_type.lower() != 's':
            print('\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
            print("Not all necessary configuration files are present. Please run 'sshyp tweak'!")
            print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n')
        s_exit()
else:
    try:
        tweak()
        s_exit()
    except KeyboardInterrupt:
        print('\n')
        s_exit()

# run function based on argument
error = 0  # create an error flag to determine if syncing should occur at end
if argument.startswith('/'):
    read_shortcut()
elif argument == '':
    try:
        no_arg()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'help' or argument == '--help' or argument == '-h':
    helper()
elif argument == 'license':
    show_license()
elif argument == 'version' or argument == '-v':
    version()
elif argument == 'add':
    add_info()
elif argument == 'add note' or argument == 'add -n' or argument == 'add password' or argument == 'add -p':
    try:
        add_entry()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'add folder' or argument == 'add -f':
    try:
        add_folder()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'edit':
    edit_info()
elif argument == 'edit rename' or argument == 'edit relocate' or argument == 'edit -r':
    try:
        rename()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'edit username' or argument == 'edit -u' or argument == 'edit password' or argument == 'edit -p' or \
        argument == 'edit url' or argument == 'edit -l' or argument == 'edit note' or argument == 'edit -n':
    try:
        edit()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'gen' or argument == 'gen update' or argument == 'gen -u':
    try:
        gen()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'copy':
    copy_info()
elif argument == 'copy username' or argument == 'copy -u' or argument == 'copy password' or argument == 'copy -p' or \
        argument == 'copy url' or argument == 'copy -l' or argument == 'copy note' or argument == 'copy -n':
    try:
        copy_data()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'shear' or argument == '-rm':
    try:
        remove_data()
    except KeyboardInterrupt:
        print('\n')
        s_exit()
elif argument == 'sync' or argument == '-s':
    sync()
else:
    error = 1  # set error flag to 1 to stop syncing
    print("\nArgument error! Please run 'sshyp help' for a list of usable commands.\n")
    s_exit()

# sync at end of program if any changes were made
if argument != '' and argument != 'help' and argument != '--help' and argument != '-h' and argument != 'license' and \
        argument != 'add' and argument != 'sync' and argument != 'edit' and argument != 'copy username' and argument \
        != 'copy -u' and argument != 'copy password' and argument != 'copy -p' and argument != 'copy url' and argument \
        != 'copy -l' and argument != 'copy note' and argument != 'copy -n' and argument != 'copy' and argument != \
        'shear' and argument != '-rm' and argument != 'version' and argument != '-v' and error == 0 and \
        argument.startswith('/') is False:
    sync()