#!/usr/bin/env python
# encoding:utf8
"""Read and copy Keepass database entries using dmenu or rofi

"""
from datetime import datetime, timedelta
import errno
from io import BytesIO
import itertools
import locale
import os
from os.path import exists, expanduser
import shlex
import sys
from subprocess import Popen, PIPE
import gpgme
from pykeepass import PyKeePass
from pykeyboard import PyKeyboard

try:
    import cPickle as pickle
except ImportError:
    import pickle
try:
    import configparser as configparser
except ImportError:
    import ConfigParser as configparser


GPG = gpgme.Context()

if sys.version_info.major < 3:
    str = unicode


def process_config():
    """Set global variables. Read the config file. Create default config file if
    one doesn't exist.

    """
    global CACHE_EXPIRY_FILE, \
           CACHE_FILE, \
           CACHE_PERIOD, \
           CACHE_PERIOD_DEFAULT, \
           CONF, \
           CONF_FILE, \
           DMENU_LEN, \
           ENV, \
           ENC
    ENV = os.environ.copy()
    ENV['LC_ALL'] = 'C'
    ENC = locale.getpreferredencoding()
    CACHE_FILE = expanduser("~/.cache/keepmenu")
    CACHE_EXPIRY_FILE = expanduser("~/.cache/keepmenu_exp")
    CACHE_PERIOD_DEFAULT = 6
    CONF_FILE = expanduser("~/.config/keepmenu/config.ini")
    CONF = configparser.ConfigParser()
    if not exists(CONF_FILE):
        with open(CONF_FILE, 'w') as conf_file:
            CONF.add_section('dmenu')
            CONF.set('dmenu', 'dmenu_command', 'dmenu')
            CONF.add_section('dmenu_passphrase')
            CONF.set('dmenu_passphrase', 'nf', '#222222')
            CONF.set('dmenu_passphrase', 'nb', '#222222')
            CONF.set('dmenu_passphrase', 'rofi_obscure', 'True')
            CONF.add_section('database')
            CONF.set('database', 'database_1', '')
            CONF.set('database', 'keyfile_1', '')
            CONF.set('database', 'gpg_key', '')
            CONF.set('database', 'pw_cache_period', str(CACHE_PERIOD_DEFAULT))
            CONF.write(conf_file)
    CONF.read(CONF_FILE)
    if CONF.has_option("database", "pw_cache_period"):
        CACHE_PERIOD = int(CONF.get("database", "pw_cache_period"))
    else:
        CACHE_PERIOD = CACHE_PERIOD_DEFAULT
    if CONF.has_option("dmenu", "l"):
        DMENU_LEN = int(CONF.get("dmenu", "l"))
    else:
        DMENU_LEN = 24


def dmenu_cmd(num_lines, prompt="Entries"):  # pylint: disable=too-many-branches
    """Parse config.ini for dmenu options

    Args: args - num_lines: number of lines to display
                 promp: prompt to show
    Returns: command invocation (as a list of strings) for
                dmenu -l <num_lines> -p <prompt> -i ...

    """
    dmenu_command = "dmenu"
    if not CONF.has_section('dmenu'):
        dmenu = [dmenu_command, "-i", "-l",
                 str(min(DMENU_LEN, num_lines)), "-p", str(prompt)]
    else:
        args = CONF.items('dmenu')
        args_dict = dict(args)
        dmenu_args = []
        if "dmenu_command" in args_dict:
            command = shlex.split(args_dict["dmenu_command"])
            dmenu_command = command[0]
            dmenu_args = command[1:]
            del args_dict["dmenu_command"]
        if "rofi" in dmenu_command:
            lines = "-i -dmenu -lines"
            # rofi doesn't support 0 length line, it requires at least -lines=1
            # see https://github.com/DaveDavenport/rofi/issues/252
            num_lines = num_lines or 1
        else:
            lines = "-i -l"
        if "l" in args_dict:
            # rofi doesn't support 0 length line, it requires at least -lines=1
            # see https://github.com/DaveDavenport/rofi/issues/252
            if "rofi" in dmenu_command:
                args_dict['l'] = min(num_lines, int(args_dict['l'])) or 1
            lines = "{} {}".format(lines, args_dict['l'])
            del args_dict['l']
        else:
            lines = "{} {}".format(lines, num_lines)
        if "pinentry" in args_dict:
            del args_dict["pinentry"]
    if prompt == "Passphrase":
        if CONF.has_section('dmenu_passphrase'):
            args = CONF.items('dmenu_passphrase')
            args_dict.update(args)
        rofi_obscure = True
        if CONF.has_option('dmenu_passphrase', 'rofi_obscure'):
            rofi_obscure = CONF.getboolean('dmenu_passphrase', 'rofi_obscure')
            del args_dict["rofi_obscure"]
        if rofi_obscure is True and "rofi" in dmenu_command:
            dmenu_args.extend(["-password"])
    extras = (["-" + str(k), str(v)] for (k, v) in args_dict.items())
    dmenu = [dmenu_command, "-p", str(prompt)]
    dmenu.extend(dmenu_args)
    dmenu += list(itertools.chain.from_iterable(extras))
    dmenu[1:1] = lines.split()
    dmenu = list(filter(None, dmenu))  # Remove empty list elements
    return dmenu


def dmenu_err(prompt):
    """Pops up a dmenu prompt with an error message

    """
    Popen(dmenu_cmd(1, prompt), stdin=PIPE, stdout=PIPE,
          env=ENV).communicate(input='')
    return


def get_databases():
    """Return all available databases from config.ini

    Returns: [(database name, keyfile, passphrase), (database2, kf, pw), ...]
    """
    args = CONF.items('database')
    args_dict = dict(args)
    dbases = [i for i in args_dict if i.startswith('database')]
    dbs = []
    for dbase in dbases:
        dbn = expanduser(args_dict[dbase])
        idx = dbase.rsplit('_', 1)[-1]
        try:
            keyfile = expanduser(args_dict['keyfile_{}'.format(idx)])
        except KeyError:
            keyfile = ''
        try:
            passw = args_dict['password_{}'.format(idx)]
        except KeyError:
            passw = ''
        if dbn:
            dbs.append((dbn, keyfile, passw))
    if not dbs:
        dmenu_err("No databases defined in {}".format(CONF_FILE))
        get_initial_db()
        dbs = get_databases()
    return dbs


def get_initial_db():
    """Ask for initial database name and keyfile if not entered in config file

    """
    db_name = Popen(dmenu_cmd(0, "Enter path to existing"
                                 "Keepass database. ~/ for $HOME is ok"),
                    stdin=PIPE,
                    stdout=PIPE).communicate()[0].decode(ENC).rstrip('\n')
    if not db_name:
        dmenu_err("No database entered. Set in ~/.config/keepmenu/config.ini")
        sys.exit()
    keyfile_name = Popen(dmenu_cmd(0, "Enter path to keyfile. "
                                   "~/ for $HOME is ok"),
                         stdin=PIPE,
                         stdout=PIPE).communicate()[0].decode(ENC).rstrip('\n')
    with open(CONF_FILE, 'w') as conf_file:
        CONF.set('database', 'database_1', db_name)
        if keyfile_name:
            CONF.set('database', 'keyfile_1', keyfile_name)
        CONF.write(conf_file)


def open_database(dbo):
    """Open keepass database and return the PyKeePass object

        Args: dbo: tuple (db path, keyfile path, password)
        Returns: PyKeePass object

    """
    dbf, keyfile, password = dbo
    if not password:
        password = get_passphrase()
    try:
        kpo = PyKeePass(dbf, password, keyfile=keyfile)
    except (IOError, OSError, IndexError) as e:  ## pylint: disable=invalid-name
        if e.args[0] == "Master key invalid." or e.args[0] == "No credentials found.":
            dmenu_err("Invalid Password or keyfile")
            # Invalidate cached pw if keyfile/password are wrong
            with open(CACHE_EXPIRY_FILE, 'wb') as cef:
                pickle.dump(datetime.now(), cef, -1)
        elif e.errno == errno.ENOENT:
            dmenu_err("Database does not exist")
        sys.exit()
    if CONF.has_option("database", "gpg_key"):
        cache_passphrase(password)
    return kpo


def get_cached_passphrase():
    """Retrieve cached passphrase if possible

    """
    # First check if the cache file exists and if it is expired
    try:
        with open(CACHE_EXPIRY_FILE, 'rb') as cef:
            expiry = pickle.load(cef)
    except IOError:
        return ""
    if datetime.now() > expiry:
        return ""
    with BytesIO() as data:
        try:
            with open(CACHE_FILE, 'rb') as cache:
                GPG.decrypt(cache, data)
                password = data.getvalue().decode(ENC).rstrip('\n')
        except IOError:
            return ""
        except gpgme.GpgmeError:
            return ""
    return password


def cache_passphrase(password):
    """Save passphrase to cache file

    """
    key = CONF.get("database", "gpg_key")
    try:
        gpg_key = GPG.get_key(key)
    except gpgme.GpgmeError:
        dmenu_err("Gpg error...invalid key. Unable to cache passphrase.")
        return
    with BytesIO(password.encode(ENC)) as data:
        try:
            with open(CACHE_FILE, 'wb') as cache:
                GPG.encrypt([gpg_key], 0, data, cache)
        except gpgme.GpgmeError:
            dmenu_err("Gpg error...unable to cache passphrase")
    with open(CACHE_EXPIRY_FILE, 'wb') as cef:
        pickle.dump(datetime.now() + timedelta(hours=CACHE_PERIOD), cef, -1)


def get_passphrase():
    """Get a database password from dmenu, pinentry or cached gpg encrypted file

    Returns: string

    """
    if CONF.has_option("database", "gpg_key"):
        password = get_cached_passphrase()
        if password:
            return password
    pinentry = None
    if CONF.has_option("dmenu", "pinentry"):
        pinentry = CONF.get("dmenu", "pinentry")
    if pinentry:
        password = ""
        out = Popen(pinentry,
                    stdout=PIPE,
                    stdin=PIPE).communicate( \
                            input=b'setdesc Enter database password\ngetpin\n')[0]
        if out:
            res = out.decode(ENC).split("\n")[2]
            if res.startswith("D "):
                password = res.split("D ")[1]
    else:
        password = Popen(dmenu_cmd(0, "Passphrase"),
                         stdin=PIPE,
                         stdout=PIPE).communicate()[0].decode(ENC).rstrip('\n')
    return password


def type_entry(entry):
    """Use PyUserInput to type the selected entry username and/or password and
    then 'Enter'.

    """
    kbd = PyKeyboard()
    if entry.username:
        kbd.type_string(entry.username)
        if entry.password:
            kbd.tap_key(kbd.tab_key)
    if entry.password:
        kbd.type_string(entry.password)
    # Not sure why we need n=2, but only seems to work that way
    kbd.tap_key(kbd.enter_key, n=2)


def type_text(data):
    """Use PyUserInput to type the given text data

    """
    kbd = PyKeyboard()
    kbd.type_string(data)


def view_all_entries(options, kp_entries):
    """Generate numbered list of all Keepass entries and open with dmenu.

    Returns: dmenu selection

    """
    num_align = len(str(len(kp_entries)))
    kp_entry_pattern = "{:>{na}} - {} - {} - {}"  # Path,username,url
    # Have to number each entry to capture duplicates correctly
    kp_entries_b = "\n".join([kp_entry_pattern.format(j, i.path, i.username, i.url, na=num_align)
                              for j, i in enumerate(kp_entries)]).encode(ENC)
    if options:
        options_b = ("\n".join(options) + "\n").encode(ENC)
        entries_b = options_b + kp_entries_b
    else:
        entries_b = kp_entries_b
    return Popen(dmenu_cmd(min(DMENU_LEN, len(options) + len(kp_entries))),
                 stdin=PIPE, stdout=PIPE, env=ENV).communicate(input=entries_b)[0].decode(ENC).rstrip()


def view_entry(kp_entry):
    """Show title, username, password, url and notes for an entry.

    Returns: dmenu selection

    """
    fields = [kp_entry.path or "Title: None",
              kp_entry.username or "Username: None",
              '**********' if kp_entry.password else "Password: None",
              kp_entry.url or "URL: None",
              "Notes: <Enter to view>" if kp_entry.notes else "Notes: None"]
    kp_entries_b = "\n".join(fields).encode(ENC)
    sel = Popen(dmenu_cmd(len(fields)), stdin=PIPE, stdout=PIPE,
                env=ENV).communicate(input=kp_entries_b)[0].decode(ENC).rstrip()
    if sel == "Notes: <Enter to view>":
        sel = view_notes(kp_entry.notes)
    elif sel == "Notes: None":
        sel = ""
    elif sel == '**********':
        sel = kp_entry.password
    return sel


def view_notes(notes):
    """View the 'Notes' field line-by-line within dmenu.

    Returns: text of the selected line for typing

    """
    notes_l = notes.split('\n')
    notes_b = "\n".join(notes_l).encode(ENC)
    sel = Popen(dmenu_cmd(min(DMENU_LEN, len(notes_l))), stdin=PIPE, stdout=PIPE,
                env=ENV).communicate(input=notes_b)[0].decode(ENC).rstrip()
    return sel


def run():
    """Main script entrypoint"""
    databases = get_databases()
    if len(databases) == 1:
        keepass = open_database(databases[0])
    else:
        inp_bytes = "\n".join(i[0] for i in databases).encode(ENC)
        sel = Popen(dmenu_cmd(len(databases), "Select Database"),
                    stdin=PIPE,
                    stdout=PIPE,
                    env=ENV).communicate(input=inp_bytes)[0].decode(ENC)
        if not sel.rstrip():
            sys.exit()
        keepass = open_database([i for i in databases if i[0] == sel.strip()][0])
    kp_entries = keepass.entries
    options = ['View/Type Individual entries']
    sel = view_all_entries(options, kp_entries)
    if not sel:
        sys.exit()
    if sel in options:
        options = []
        sel = view_all_entries(options, kp_entries)
        entry = kp_entries[int(sel.split('-', 1)[0])]
        text = view_entry(entry)
        type_text(text)
    else:
        try:
            entry = kp_entries[int(sel.split('-', 1)[0])]
        except ValueError:
            sys.exit()
        type_entry(entry)


if __name__ == '__main__':
    process_config()
    run()

# vim: set et ts=4 sw=4 :
