#!/usr/bin/env python
# vim: set ts=8 sw=4 sts=4 et ai tw=79:
"""
pstore -- Python Password Protected Store
Copyright (C) 2010,2012,2013  Walter Doekes <wdoekes>, OSSO B.V.

    This application is free software; you can redistribute it and/or modify
    it under the terms of the GNU Lesser General Public License as published
    by the Free Software Foundation; either version 3 of the License, or (at
    your option) any later version.

    This application is distributed in the hope that it will be useful, but
    WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this application; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307,
    USA.
"""
from __future__ import absolute_import

import os
import sys
from getopt import GetoptError, gnu_getopt
from getpass import getpass, getuser
from tempfile import mktemp

from pstorelib import VERSION_STRING
from pstorelib.bytes import BytesIO, sendfile
from pstorelib.crypt import PStoreCrypt, CryptoWriter
from pstorelib.exceptions import (PStoreException, UserError, NotAllowed,
                                  NotFound)
from pstorelib.server import Backend


# Optional. If they don't exist, you won't notice their features.
SSH_ASKPASS_APP = '/usr/local/bin/ssh_echopass'
SSH_ASKPASS_PRELOAD = '/usr/local/lib/ssh_askpass.so'


class PStoreInterface(object):
    PASSWORD_PROPERTY = 'password'

    def __init__(self, config):
        self.config = config
        self.backend = Backend(urls=config['store_urls'],
                               user=config['user'],
                               verbose=config['verbose'])

    ###########################################################################
    ## SHORTCUTS
    ###########################################################################

    def add_machine(self, machine, new_password, allowed_users):
        if new_password in (None, ''):
            raise UserError('empty password rejected')
        fp = BytesIO(new_password)
        self.propset(machine, self.PASSWORD_PROPERTY, fp,
                     allowed_users=allowed_users)

    def change_password(self, machine, new_password):
        if new_password in (None, ''):
            raise UserError('empty password rejected')
        fp = BytesIO(new_password)
        self.propset(machine, self.PASSWORD_PROPERTY, fp,
                     allowed_users=True)  # <-- True means "use existing users"

    def change_permissions(self, objectid, allow, revoke):
        object = self.backend.get_object(objectid)

        users = object['users'].keys()
        # It doesn't make sense to check these on the server. The only
        # server-side check we do is on the resulting new_users list.
        if not all(i in users for i in revoke):
            raise UserError('one or more users cannot be revoked')
        if any(i in users for i in allow):
            raise UserError('one or more users already allowed')

        new_users = list(set(users).difference(set(revoke)).union(set(allow)))
        self.alter_permissions(objectid, new_users)

    def get_objects(self, objectid_contains, allowed_only, verbose):
        return self.backend.get_objects(objectid_contains=objectid_contains,
                                        allowed_only=allowed_only,
                                        verbose=verbose)

    def get_password(self, identifier):
        # Get object.
        object = self.backend.get_object(identifier)
        # The object fetched is already quite verbose. The password might be
        # in there already.
        property = self.PASSWORD_PROPERTY
        try:
            info = object['properties'][property]     # KeyError if not
            decrypt_with = info['data'].decrypt_with  # AttributeError if None
        except (AttributeError, KeyError):
            fp = self.propget(identifier, property)
        else:
            if info['enctype'] == 'none':
                print >>sys.stderr, 'WARNING: password as public property?'
            fp = decrypt_with(None)

        password = fp.read()
        return password, object

    def get_properties(self, objectid, allowed_only=True):
        '''Get the available properties, possibly with contents.'''
        object = self.backend.get_object(objectid, allowed_only=allowed_only)
        return object['properties']

    ###########################################################################
    ## REGULAR INTERFACE
    ###########################################################################

    def alter_permissions(self, objectid, allowed_users):
        # 1. Get user info for the new list of allowed users.
        # (may raise NotFound, NotAllowed, NoNonce)
        keys = self.backend.get_keys(allowed_users)
        assert len(keys) == len(allowed_users)

        # 2. Get all *encrypted* properties.
        # FIXME: get_full_object ?
        object = self.backend.get_object(objectid)
        properties = {}
        for property, info in object['properties'].items():
            if info['enctype'] == 'none':
                continue
            property = str(property)  # force bytestring; using it as dict key
            properties[property] = self.propget(objectid, property)

        # 2a.
        # FIXME: there should be a way to denote that we are not updating the
        # properties for a particular user.
        # that would allow us to only send a new list of usernames when we only
        # want to revoke; and to only send property contents for the new users.

        # 3. Re-encrypt those properties according to the new allowed users
        # list.
        files = []
        for property, fp in properties.items():
            writer = CryptoWriter(fp=fp)
            if 'PSTORE_NOINPUT' not in os.environ:
                print >>sys.stderr, 'INFO: encrypting..'
            for user in allowed_users:
                encryptedfp = writer.encrypt_with(keys[user])
                files.append((property, user, encryptedfp))

        # 4. Upload the whole batch in a single go, along with the user
        # modifications.
        self.backend.propupd(objectid, files=files)

    def propget(self, objectid, property):
        return self.backend.propget(objectid, property)

    def propset(self, objectid, property, fp, allowed_users=None):
        """The ``allowed_users`` argument can be one of three things:
        - ``None`` which means the property is public (unencrypted)
        - ``True`` which means: take the user list from the server. This won't
          work when creating a new objectid.
        - An explicit list with usernames.

        If allowed_users is None, the property is public (unencrypted). If
        allowed_users is True, the object should already exist: we take the
        current list of allowed users. If the object doesn't exist,
        allowed_users should be a list of usernames.
        """
        if allowed_users:
            if allowed_users is True:
                # Fetch the currently allowed users list.
                data = self.backend.get_object(objectid)
                # Are we allowed to do anything?
                if self.config['user'] not in data['users'].keys():
                    # This will fail.. but we'll let the server do this check.
                    # That's where the check should happen anyway.
                    print >>sys.stderr, 'WARNING: this will fail'
                # Set the allowed users list.
                allowed_users = data['users']
                del data

            # May raise UserError
            keys = self.backend.get_keys(allowed_users)
            assert len(keys) == len(allowed_users)

            # Encrypt the value.
            files = []
            writer = CryptoWriter(fp=fp)
            if 'PSTORE_NOINPUT' not in os.environ:
                print >>sys.stderr, 'INFO: encrypting..'
            for user in allowed_users:
                encryptedfp = writer.encrypt_with(keys[user])
                files.append((property, user, encryptedfp))
        else:
            # The value is public.
            files = [(property, '*', fp)]

        # FIXME: rename "files" to "values" or something
        self.backend.propset(objectid, property, files=files)

    def validate(self):
        return self.backend.validate()


def parse_pstorerc():
    '''
    Read ~/.pstorerc for default arguments.

    "There was a facility that would execute a bunch of commands stored
    in a file; it was called runcom for "run commands", and the file
    began to be called "a runcom". rc in Unix is a fossil from that
    usage."
    '''
    # TODO: emit warning if .pstorerc is found but not readable
    rc_args, file, lines = [], None, None
    try:
        file = open(os.environ['HOME'] + '/.pstorerc', 'r')
        rc_args.extend(file.read().split())
    except:
        pass
    finally:
        if file:
            file.close()
    return rc_args


def parse_options(args):
    '''
    Parse command line options, with default options read from the
    .pstorerc file prefixed. This way you can set defaults there, like
    the --store-url, but it can be overridden manually.
    '''
    rc_args = parse_pstorerc()

    # Quick hack to warn people of duplicate store-urls.
    has_default_store_url = any(i.startswith('--store-url') for i in rc_args)
    has_new_store_url = any(i.startswith('--store-url') for i in args[1:])
    if has_default_store_url and has_new_store_url:
        msg = 'WARNING: setting --store-url does not override .pstorerc'
        print >>sys.stderr, msg  # FIXME?

    try:
        optlist, args = gnu_getopt(
            [args[0]] + rc_args + args[1:],  # inject defaults
            'hcPp:' + 'ak:l:u:v',  # command options + configuration options
            ('help', 'create', 'chpass', 'lookup-machine', 'propedit',
             'consistency-check') +
            ('all', 'private-key=', 'store-url=', 'user=', 'verbose')
        )
    except GetoptError, e:
        raise UserError(unicode(e))

    config = {
        'show_all': False,
        'store_urls': [],
        'user': getuser(),
        'verbose': False
    }
    if 'HOME' in os.environ:
        config['privkey_file'] = '%s/.ssh/id_rsa' % os.environ['HOME']

    command = None
    for option, arg in optlist:
        # Command options
        if option in ('--help', '-h'):
            command = 'help'  # always allow -h
        elif option in ('--create', '-c',
                        '--change-password', '-P',
                        '--propedit', '-p',
                        '--lookup-machine',
                        '--consistency-check'):
            if not command:
                if option in ('--create', '-c'):
                    command = 'create'
                elif option in ('--change-password', '-P'):
                    command = 'chpass'
                elif option in ('--lookup-machine',):
                    command = 'list_minimal'
                elif option in ('--propedit', '-p'):
                    if arg == 'l':
                        command = 'proplist'
                    elif arg == 'g':
                        command = 'propget'
                    elif arg == 's':
                        command = 'propset'
                    elif arg == 'e':
                        command = 'propsetenc'
                    else:
                        assert False
                elif option in ('--consistency-check',):
                    command = 'check_consistency'
                else:
                    assert False
            elif command != 'help':
                command = 'multiple_commands'

        # Configuration options
        elif option in ('--all', '-a'):
            config['show_all'] = True
        elif option in ('--private-key', '-k'):
            config['privkey_file'] = arg
        elif option in ('--store-url',):
            while len(arg) and arg[-1] == '/':  # strip trailing slashes
                arg = arg[0:-1]
            if (not arg.startswith('http://') and
                not arg.startswith('https://')):
                raise UserError('bad store-url; we only do http(s)', arg)
            config['store_urls'].append(arg)
        elif option in ('--user', '-u'):
            config['user'] = arg
        elif option in ('--verbose', '-v'):
            config['verbose'] = True

        # Evil hacks to allow the password input into ssh directly
        elif option in ('-l',):
            # XXX: move this to the ssh_host_hack bit..? or rename ssh_host to
            # ssh_arg or something..
            # Assume args[1] is the machine name here
            machine_name = args[1]
            # "x" (or "-") => use machine_name
            if arg in ('-', 'x') and len(args) > 1:
                config['ssh_host'] = machine_name
            # "username" => use username@machine_name
            elif '.' not in arg:
                config['ssh_host'] = arg + '@' + machine_name
            # "host-part." => use host-part.machine-domain"
            elif arg.endswith('.'):
                config['ssh_host'] = arg + machine_name.split('.', 1)[-1]
            # "whatever-else" => use "whatever-else"
            else:
                config['ssh_host'] = arg
            del machine_name
            print 'Attempting ssh login with arg:', config['ssh_host']
            print ('(if you see "Host key verification failed", you need to '
                   'ssh manually)')

        else:
            raise NotImplementedError('Unhandled option', option)

    if command == 'multiple_commands':
        raise UserError('multiple command options encountered, see -h')

    return command, args[1:], config  # drop argv0


def parse_allow_revoke(args):
    '''
    Expects a list of usernames with + or - prefix. Split these
    usernames into two lists and return. Note that everything is in
    lower case.
    '''
    allow, revoke = [], []
    for arg in args:
        if len(arg) < 2 or arg[0] not in ('+', '^'):
            print >>sys.stderr, ('WARNING: skipping unparseable allow/revoke '
                                 'statement %s' % (arg,))
        elif arg[0] == '+':
            allow.append(arg[1:].lower())
        elif arg[0] == '^':
            revoke.append(arg[1:].lower())
        else:
            assert False
    return allow, revoke


def run_command(command, args, config):
    if config['verbose']:
        print >>sys.stderr, ('- Executing "%s" on stores %r as user "%s"' %
                             (command, config['store_urls'], config['user']))

    if command == 'check_consistency':
        pass
    elif args:
        name = args.pop(0).lower()
    elif command == 'default':
        name = ''  # lookup all
    else:
        raise UserError('need machine name')

    if command in ('check_consistency', 'proplist'):
        if len(args):
            raise UserError('unexpected argument')
    elif command in ('propget', 'propset', 'propsetenc'):
        if len(args) != 1:
            raise UserError('invalid number of arguments')
        property = args.pop(0).lower()
    else:
        allow, revoke = parse_allow_revoke(args)

    if command == 'default':
        if not allow and not revoke:
            command = 'show_or_list'
        else:
            command = 'allow_and_revoke'

    if command in ('create', 'chpass'):
        if command == 'create' and revoke:
            raise UserError('cannot revoke users when adding a new machine')

        if 'PSTORE_NOINPUT' in os.environ:
            # For tests we'd like to create passwords automatically.
            new_password = 'example-password'
        else:
            new_password = getpass('Type new machine password: ')
            new_password2 = getpass('Type new machine password again: ')
            if new_password != new_password2:
                raise UserError('passwords do not match')
            del new_password2

    pstore = PStoreInterface(config)

    if command == 'create':
        pstore.add_machine(name, new_password, [config['user']] + allow)

    elif command == 'chpass':
        pstore.change_password(name, new_password)
        if allow or revoke:
            pstore.change_permissions(name, allow, revoke)

    elif command == 'allow_and_revoke':
        if not allow and not revoke:
            raise UserError('need something to allow/revoke')
        pstore.change_permissions(name, allow, revoke)

    elif command == 'list_minimal':
        # Do not ask for passwords when called from a bash_completion script.
        PStoreCrypt.get().askpass = PStoreCrypt.ASKPASS_NEVER
        allowed_only = (not config['show_all'])
        objects = pstore.get_objects(objectid_contains=name,
                                     allowed_only=allowed_only,
                                     verbose=False)
        for objectid in sorted(objects.keys()):
            print objectid

    elif command == 'propget':
        # TODO: Add the possibility to replace stdout with a file..
        fp = pstore.propget(name, property)
        # Use a sendfile(2) compatible call to push everything from the fp to
        # stdout.
        sendfile(sys.stdout, fp)

    elif command == 'proplist':
        allowed_only = (not config['show_all'])
        properties = pstore.get_properties(name,
                                           allowed_only=allowed_only)
        if config['verbose']:
            for property in sorted(properties.keys()):
                value = properties[property]
                info = 'size=%d' % (value['size'],)
                if value['enctype'] == 'none':
                    info += ', unencrypted'
                print '%s (%s)' % (property, info)
        else:
            for property in sorted(properties.keys()):
                print property

    elif command in ('propset', 'propsetenc'):
        # TODO: Add the possibility to replace stdin with a file..
        fp = sys.stdin
        # Warn when storing unencrypted properties.
        if command == 'propset':
            if 'PSTORE_NOINPUT' not in os.environ:
                print >>sys.stderr, 'NOTICE: not encrypting the value'
        # Set allowed users automatically.
        allowed_users = (None, True)[command == 'propsetenc']
        pstore.propset(name, property, fp, allowed_users=allowed_users)

    elif command == 'show_or_list':
        # Try show first. If that fails, do a listing.
        found_machine = False

        if name:
            try:
                machine_password, machine_info = pstore.get_password(name)
            except (NotAllowed, NotFound):
                # A superadmin gets a 404 on no-such-object. A mere mortal gets
                # a 403 instead. Catch both.
                pass
            else:
                run_show_password(machine_info, machine_password,
                                  ssh_host_hack=config.get('ssh_host', None))
                found_machine = True

        # If there was no name, or the name wasn't a machine, do a listing.
        if not found_machine:
            limited = not config['show_all']
            machines = pstore.get_objects(objectid_contains=name,
                                          allowed_only=limited, verbose=True)
            if not machines:
                raise UserError('no visible machines match %s' % (name,))
            fmt = '%s %-25s %s'
            print fmt % (' ', 'Machine', 'User access')
            print '-' * 72
            for machine in sorted(machines.keys()):
                properties = machines[machine]
                print fmt % ('+', machine,
                             ', '.join(sorted(properties['users'].keys())))

    elif command == 'check_consistency':
        errors = pstore.validate()
        if not errors:
            if config['verbose']:
                print 'Database is consistent.'
        else:
            for error in errors:
                print 'There is a problem with %s %s' % tuple(error[0:2])
                print '  problem:', error[2]
                print '  solution:', error[3]
                print
            raise UserError('database is inconsistent')

    else:
        raise NotImplementedError('Unknown command', command)


def run_show_password(machine_info, machine_password, ssh_host_hack=None):
    # FIXME FIXME FIXME
    # machine_password should be in some kind of mutable XORed form,
    # so it doesn't easily show up in memory dumps
    # note that the memory is already exposed earlier on.. go back and fix
    # it there too.
    # FIXME FIXME FIXME

    # Er.. why do we assume that the user wants all this on his screen?
    for property in sorted(machine_info['properties'].keys()):
        if property == PStoreInterface.PASSWORD_PROPERTY:
            continue  # do that one last..
        info = machine_info['properties'][property]
        if info['enctype'] == 'none':
            if info['data'] is not None:
                print '%s = %s' % (property,
                                   info['data'].decrypt_with(None).read())
            else:
                print '%s = (%s byte document)' % (property, info['size'])
        else:
            print '%s = (%s byte encrypted)' % (property, info['size'])

    # Run ssh
    if ssh_host_hack:
        if 'DISPLAY' not in os.environ:
            os.environ['DISPLAY'] = 'whatever:0'
        os.environ['LD_PRELOAD'] = SSH_ASKPASS_PRELOAD
        os.environ['SSH_ASKPASS'] = SSH_ASKPASS_APP
        os.environ['SSH_PASSWORD'] = mktemp()
        # Not really secure.. using the filesystem.. but it'll have
        # to do for now..
        temp = open(os.environ['SSH_PASSWORD'], 'w')
        try:
            temp.write(machine_password)
        finally:
            temp.close()
            del temp
        try:
            os.execve('/usr/bin/ssh',
                      ['ssh', '-oPubkeyAuthentication=no', ssh_host_hack],
                      os.environ)
        finally:
            # We don't usually get here, but if we do, make sure we clean up
            # the password.
            temp = open(os.environ['SSH_PASSWORD'], 'w')
            temp.write('X' * 1024)
            temp.close()
            os.unlink(os.environ['SSH_PASSWORD'])
    # Show password
    else:
        if hasattr(os, 'isatty') and os.isatty(sys.stdout.fileno()):
            # ansi color the password so it's less visible
            # TODO: invert this if the console is not white-on-black
            machine_password = '\x1b[7;30m%s\x1b[0m' % machine_password
        print 'password = %s' % (machine_password,)


def run_usage(verbose=False):
    # Check if we can use the evil ssh auto-login.
    evil_ssh_hacks = os.path.exists(SSH_ASKPASS_PRELOAD)

    # Select verbosity.
    if verbose and evil_ssh_hacks:
        matcher = lambda i: i in (0, 1, 2)
    elif verbose:
        matcher = lambda i: i in (0, 1)
    else:
        matcher = lambda i: i in (0,)

    # The command options.
    commands = (
        ('  --create, -c          ', 0,
         'Add new machine (instead of looking up one)'),
        ('  --chpass, -P          ', 0,
         'Change machine password'),
        ('  --help, -h            ', 0,
         'Show help and exit (use -v for verbose help)'),
        ('  --lookup-machine      ', 1,
         'Look up machine names only (*L)'),
        ('  --propedit=A, -p A    ', 1,
         'Act on properties; list/get/set/encrypt (*P)'),
        ('  -l [x|SSH_ARG]        ', 2,
         'Evil hack to ssh(1) to machine directly (*S)'),
        ('  --consistency-check   ', 1,
         'Check database consistency (*C)'),
    )

    # The configuration options.
    configs = (
        ('  --all, -a             ', 0,
         'Privileged users use this to show more machines'),
        ('  --user=USER, -u       ', 0,
         'Your pstore username (defaults to $USER)'),
        ('  --private-key=FILE    ', 1,
         'Location of sshrsa private key (old-style)'),
        ('  --store-url=URL       ', 0,
         'Backend URL (supply more as fallbacks)'),
        ('  --verbose, -v         ', 0,
         'Verbose mode'),
    )

    # Undocumented environment variables:
    #  * PSTORE_NOINPUT
    #    If nonempty, will remove all password/value prompts. Used by the
    #    automated tests.

    # Extra stuff.
    extra = ''
    if matcher(1):
        extra += ('\n(*C) Consistency checks should be run periodically '
                  'by the admin.'
                  '\n(*L) For bash completion. Will never ask for a password.'
                  '\n(*P) Use one of -pl, -pg, -ps, -pe, where -pe is the '
                  'encrypted -ps.')
    if matcher(2):
        extra += ('\n(*S) SSH_ARG can be one of "username", '
                  '"[username@]replace-host.",\n'
                  '    "[username@]entire-host.name". '
                  'Requires ssh_askpass.so.')
    if extra:
        extra = '\n' + extra

    # Combine the data.
    commands = '\n'.join(i[0] + i[2] for i in commands if matcher(i[1]))
    configs = '\n'.join(i[0] + i[2] for i in configs if matcher(i[1]))

    # Print the help.
    print '''%(title)-60s||
                                                            ||
Usage:                                                      ||
  pstore machine                                ======||====== Python
    (look up a machine password)                      ||    || Protected
  pstore machine +user1 ^user2                        ||    || Password
    (allow (+) and revoke (^) access)                 ======|| Store
  pstore -c machine [+user1]                          ||
    (add machine and allow access)                    ||
  pstore -P machine [+user1] [^user2]                 ||
    (change machine password and allow/revoke access)
  echo Under the cupboard | pstore -ps machine location
    (add unencrypted location property to machine; use -pl to list, -pg
     to get and -pe to store *encrypted* properties)

Command options:
%(commands)s
Configuration options:
%(configs)s%(extra)s''' % {
        'commands': commands,
        'configs': configs,
        'extra': extra,
        'title': 'pstore %s CLI' % (VERSION_STRING,),
    }


def pstore(args):
    command, args, config = parse_options(args)

    if command == 'help':
        run_usage(config['verbose'])
        sys.exit(0)

    if not config['store_urls']:
        raise UserError('need at least one store-url to function')

    # If you set the PSTORE_NOINPUT environment variable, we won't ask for
    # passwords, but return username + '2' as password instead. This allows
    # automated tests to run.
    if bool('PSTORE_NOINPUT' in os.environ):
        askpass = PStoreCrypt.ASKPASS_USERNAME2     # automated tests
    else:
        askpass = PStoreCrypt.ASKPASS_DEFAULT       # default

    # Expand tilde (~) in file names.
    ssh_privkey_file = (os.path.expanduser(config.get('privkey_file', '')) or
                        None)

    # Not so nice.. but we have to do this somewhere..
    PStoreCrypt.create(askpass=askpass, sshrsa_privkey_file=ssh_privkey_file)

    run_command(command or 'default', args, config)


def main(args):
    try:
        pstore(args)
    except KeyboardInterrupt:
        sys.exit(130)  # 128+SIGINT
    except PStoreException, e:
        # TODO: get __cause__ in verbose mode?
        print >>sys.stderr, ': '.join(str(i) for i in e.args)
        sys.exit(1)


if __name__ == '__main__':
    main(sys.argv)
