#!/usr/bin/env python3

import argparse
import base64
import functools
import hashlib
import json
import os
import signal
import sqlite3
import subprocess
import urllib.request

db_path = os.path.join(os.environ['HOME'],'.passphrase')
current_schema_version = 1
#Description and arguments
def parse_args():
    parser = argparse.ArgumentParser(description='Generate a passphrase')
    subparsers = parser.add_subparsers(help='subcommands: forbid', dest='subcommand')

    forbid_parser = subparsers.add_parser('forbid', help='forbid a character from appearing in passwords')
    forbid_parser.add_argument('-site', required=True)
    forbid_parser.add_argument('-forbid', required=True)

    set_configuration_parser = subparsers.add_parser('configure', help='set configuration values')
    set_configuration_parser.add_argument('-key', required=True)
    set_configuration_parser.add_argument('-value', required=True)

    new_parser = subparsers.add_parser('new', help='add a new login', aliases=['n'])
    new_parser.add_argument('-site', required=True)
    new_parser.add_argument('-username', required=True)
    new_parser.add_argument('-length', required=False)

    generate_parser = subparsers.add_parser('generate', help='generate a passphrase', aliases=['g'])
    generate_parser.add_argument('-site', required=True)

    lookup_parser = subparsers.add_parser('lookup', help='lookup a username', aliases=['l'])
    lookup_parser.add_argument('-site', required=True)

    delete_parser = subparsers.add_parser('delete', help='delete a login', aliases=['d'])
    delete_parser.add_argument('-site', required=True)

    return parser.parse_args()

def migrate_db(conn):
    c = conn.cursor()
    c.execute('''PRAGMA user_version''')
    version = int(c.fetchone()[0])
    while True:
        if version == 0:
            conn.execute('''
CREATE TABLE IF NOT EXISTS settings (
site text UNIQUE NOT NULL,
username text UNIQUE NOT NULL,
length integer)''')
            conn.execute('''
CREATE TABLE IF NOT EXISTS forbidden_characters (
site text UNIQUE NOT NULL,
forbidden text NOT NULL,
UNIQUE(site, forbidden) ON CONFLICT ROLLBACK)''')
        elif version == 1:
            conn.execute('''
ALTER TABLE settings RENAME TO logins''')
            conn.execute('''
CREATE TABLE configurations (
key text UNIQUE NOT NULL,
value text NOT NULL)''')
            conn.execute('''
INSERT INTO configurations VALUES ("pinentry", "pinentry")''')
        else:
            c.execute('''PRAGMA user_version = ''' + str(version))
            conn.commit()
            break
        version += 1

def db_setup():
    conn = sqlite3.connect(db_path)
    migrate_db(conn)
    return conn

class TimeoutExpired(BaseException):
    pass

def timeout_handler(signum, frame):
    raise TimeoutExpired()

class PinEntry:
    """Interact with the PinEntry front-end"""
    process = None
    default_timeout = 30

    def __init__(self, pinentry, default_timeout = 300):
        proc = subprocess.run(["tty"], stdout=subprocess.PIPE, universal_newlines=True)
        tty = proc.stdout.strip()
        self.default_timeout = default_timeout
        self.process = subprocess.Popen([pinentry, "--display", os.environ["DISPLAY"]], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        output = self.communicate()
        if output[0:2] != "OK":
            raise Exception("pinentry didn't initialize correctly: " + output)
        output = self.communicate("OPTION ttyname="+tty)
        self.set_time_out(self.default_timeout)

    def communicate(self, message=None, timeout=0):
        if timeout == 0:
            timeout = self.default_timeout
        signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(timeout + 1)
        output = ""
        try:
            if message != None:
                self.process.stdin.write(message + "\n")
                self.process.stdin.flush()
            output = self.process.stdout.readline()
        except TimeoutExpired as e:
            _, errors = self.process.communicate(timeout=10)
            raise Exception(str(errors))
        except Exception as e:
            signal.alarm(0)
            signal.signal(signal.SIGALRM, signal.SIG_DFL)
            raise e
        signal.alarm(0)
        signal.signal(signal.SIGALRM, signal.SIG_DFL)
        return output

    def check_result(self, result):
        if result[0:2] != "OK":
            raise Exception("unexpected pinentry response: " + result)

    def set_time_out(self, timeout):
        message = "SETTIMEOUT " + str(timeout)
        self.check_result(self.communicate(message=message))

    def set_descriptive_text(self, text):
        message = "SETDESC " + str(text)
        self.check_result(self.communicate(message=message))

    def set_prompt(self, prompt):
        message = "SETPROMPT " + str(prompt)
        self.check_result(self.communicate(message=message))

    def get_pin(self):
        pin = self.communicate(message="GETPIN", timeout=300)
        if pin[0] == "D":
            self.check_result(self.communicate())
            return pin[2:-1]
        raise Exception("unexpected pinentry response: " + pin)

class Context:
    """Everything needed to generate a password but the secret"""

    def __init__(self, site):
        self.site = site
        self.username = None
        self.length = None
        self.forbidden_characters = []

    def insert(self):
        """Store the record"""
        req_body = json.dumps({'username': self.username,
            'length': int(self.length) if self.length != None else None,
            'forbidden_characters': self.forbidden_characters,
            })
        req = urllib.request.Request("http://localhost:8000/sites/" + self.site,method="PUT",data=bytes(req_body, encoding='utf-8'))
        status = urllib.request.urlopen(req).getcode()
        if status != 201 and status != 204:
            raise Exception("bad status from server: " + status)

    def delete(self):
        """Delete the record"""
        req = urllib.request.Request("http://localhost:8000/sites/" + self.site,method="DELETE")
        status = urllib.request.urlopen(req).getcode()
        if status == 204:
            pass
        elif status == 404:
            print("No record found for site: " + self.site)
        else:
            raise Exception("bad status from server: " + status)

    def retrieve(self):
        """Retrieve an existing record for self.site"""
        resp = None
        try:
            resp = urllib.request.urlopen("http://localhost:8000/sites/" + self.site)
        except urllib.error.HTTPError as e:
            if e.code == 404:
                raise KeyError
            else:
                raise e
        val = json.loads(resp.read())[0]
        self.username = val['username']
        self.length = val['length']
        self.forbidden_characters = val['forbidden_characters']

    def get_secret(self, conn):
        c = conn.cursor()
        c.execute('''SELECT value from configurations WHERE key = ?''', ("pinentry",))
        pinentry_program = c.fetchone()[0]

        pinentry = PinEntry(pinentry_program)
        pinentry.set_descriptive_text("Enter passphrase for " + self.username + " @ " + self.site + ":")
        pinentry.set_prompt("Passphrase:")
        return pinentry.get_pin()


    def generate(self, conn):
        secret = self.get_secret(conn)
        passwd = self.username + secret
        dk = hashlib.pbkdf2_hmac('sha256',
                                 passwd.encode('utf-8', 'ignore'),
                                 self.site.encode('utf-8', 'ignore'),
                                 500000)
        passphrase = base64.b64encode(dk).decode('utf-8')
        if self.length:
            passphrase = passphrase[:self.length]
        return passphrase

    def prune(self, passphrase):
        """Remove forbidden characters in the resultant password"""
        for seq in self.forbidden_characters:
            segments = passphrase.split(seq)
            passphrase = functools.reduce(lambda res, nxt: res + nxt, segments, "")
        return passphrase

#Main function
def main():
    args = parse_args()
    conn = db_setup()

    if args.subcommand == 'forbid':
        c = Context(args.site)
        try:
            c.retrieve()
        except KeyError:
            print('no records for ' + args.site)
            return
        c.forbidden_characters.append(args.forbid)
        c.insert()
        print('forbidding ' + args.forbid + ' on ' + args.site)
        return
    if args.subcommand == 'configure':
        conn.execute('''INSERT OR REPLACE INTO configurations VALUES (?,?)''',
                     (args.key, args.value))
        conn.commit()
        print('configured ' + args.key + ' to ' + args.value)
        return
    if args.subcommand == 'generate' or args.subcommand == 'g':
        c = Context(args.site)
        try:
            c.retrieve()
        except KeyError:
            print('no records for ' + args.site)
            return
        print(c.prune(c.generate(conn)))
    if args.subcommand == 'lookup' or args.subcommand == 'l':
        c = Context(args.site)
        try:
            c.retrieve()
            print(c.username + ' @ ' + c.site)
        except KeyError:
            print('no records for ' + args.site)
    if args.subcommand == 'new' or args.subcommand == 'n':
        c = Context(args.site)
        try:
            c.retrieve()
            if c.username != args.username:
                print("WARNING: New username provided. This will overwrite the old username (%s)" %
                      c.username)
            if c.length != args.length:
                print("WARNING: New length provided. This will overwrite the old length (%s)" %
                      (str(c.length) if c.length else 'infinite'))
        except KeyError:
            pass
        c.username = args.username
        c.length = args.length
        c.insert()

    if args.subcommand == 'delete' or args.subcommand == 'd':
        c = Context(args.site)
        try:
            c.delete()
        except KeyError:
            print('no records for ' + args.site)

if __name__ == "__main__":
    main()
 
