#!/usr/bin/python
# -*- coding: utf-8 -*-
#
#  privacyIDEA
#  (c) 2014 Cornelius Kölbel, cornelius@privacyidea.org
#
# This code is free software; you can redistribute it and/or
# modify it under the terms of the GNU AFFERO GENERAL PUBLIC LICENSE
# License as published by the Free Software Foundation; either
# version 3 of the License, or any later version.
#
# This code 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 AFFERO GENERAL PUBLIC LICENSE for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
'''
This tool is used to fetch the SSH app information from 
the privacycIDEA server and push the information via
salt to the clients. 
'''

VERSION = '1.3'
DEBUG = False
DESCRIPTION=__doc__

from getopt import getopt, GetoptError
import getpass
import argparse
from privacyideautils.clientutils import *
try:
    import salt.client
except ImportError as e:
    print e
    print("You need to install salt-common or run this program on"
          "the salt master.")
    
SALT_CONFIG = "/etc/salt/master"
USER_KEY = "option_user"
DEFAULT_USER = "root"


def help():
    print "%s --url=<url> --admin=<adminusername>" % sys.argv[0]
    print """ 
    --url/-U        : The base url of the privacyIDEA server. Something like
                      http://localhost:5000 or https://privacyidea:443
    --admin=/-a     : administrator name
    --password=/-p  : the password of the administrator
    --passwordfile= : a file, that holds the password of the administrator
    --purgekeys     : If this bool parameter is given, existing keys will
                      be purged, before the new keys from privacyIDEA will
                      be pushed to the clients.
 """
    sys.exit(2)


def create_arguments():
    parser = argparse.ArgumentParser(description=DESCRIPTION,
                                     fromfile_prefix_chars='@')
    parser.add_argument("-U", "--url",
                        help="The URL of the privacyIDEA server including "
                        "protocol and port like "
                        "https://localhost:5001",
                        required=True)
    parser.add_argument("-a", "--admin",
                        help="The name of the administrator like "
                        "admin@admin or admin",
                        required=True)
    parser.add_argument("-r", "--adminrealm",
                        help="The realm of the administrator like "
                        "'admin'",
                        default="")
    parser.add_argument("-p", "--password",
                        help="The password of the administrator. Please "
                        "avoid to post the password at the command line. "
                        "You will be asked for it - or you can provide the "
                        "password in a configuration file. "
                        "Note, that you can write a file password.txt "
                        "containing two lines '--password' and the second "
                        "line the password itself and add this to the "
                        "command line with @password.txt")
    parser.add_argument("-v", "--version",
                        help="Print the version of the program.",
                        action='version', version='%(prog)s ' + VERSION)
    parser.add_argument("--purgekeys",
                        help="If this bool parameter is given, existing keys "
                        "will be purged, before the new keys from privacyIDEA "
                        "will be pushed to the clients.",
                        action="store_true")
    args = parser.parse_args()
    return args

def get_machine_user_dict(machines):
    '''
    returns a dictionary with machines as keys and a list of users as values. 
    '''
    machine_user = {}
    for machine_id in machines.keys():
        machine = machines[machine_id]
        machinename = machine.get("machinename")
        if machinename:
            user = machine.get("options", {}).get(USER_KEY, DEFAULT_USER)
            if machinename not in machine_user.keys():
                # The machine does not exist, yet                
                machine_user[machinename] = set()
            machine_user[machinename].add(user)
    return machine_user

def ssh_purge_keys(minion, users):
    '''
    Purge the SSH keys from all users on the minion
    '''
    salt_client = salt.client.LocalClient(c_path=SALT_CONFIG)
    print "Purge keys on minion %s" % minion
    for user in users:
        print " * for user %s" % user
        res = salt_client.cmd(minion,
                              "ssh.auth_keys",
                              kwarg={"user": user})
        for key in res.get(minion).keys():
            res = salt_client.cmd(minion,
                                  "ssh.rm_auth_key", 
                                  kwarg={"user": user,
                                         "key": key})
            print " * %s" % res
    

def set_ssh(machine,
            purgekeys=False):
    '''
    :param machine: dictionary with keys like 
            ip,
            is_active,
            options,
            auth_item,
            machine_name
           
    
            import salt.client
            local = salt.client.LocalClient(c_path='/etc/salt/master') 
            local.cmd('salt-minion',
                      'ssh.rm_auth_key',
                      kwarg={"user": "cornelius", "key": "aaa1234"})
    '''
    # We assume, we run on the master
    salt_client = salt.client.LocalClient(c_path=SALT_CONFIG)
    
    minion = machine.get("machinename")
    user = machine.get("options", {}).get(USER_KEY, DEFAULT_USER)
    authitem = machine.get("auth_item")
    ssh_key = authitem.get("sshkey")
    
    print("===============================================")
    print("processing minion %s" % minion)
    
    key_elements = ssh_key.split()
    if len(key_elements) == 1:
        client_key = key
        client_comment = "unknown key"
        client_enc = "ssh-rsa"
    elif len(key_elements) == 2:
        client_key = key_elements[1]
        client_enc = key_elements[0]
        client_comment = "unknown key"
    elif len(key_elements) == 3:
        client_key = key_elements[1]
        client_enc = key_elements[0]
        client_comment = key_elements[2]
    else:
        print("Could not handle the key for minion %s."
              " Could not split key" % minion)
        return False
    params = {"user": user,
              "key": client_key}

    if machine.get("is_active", True) is False:
        # remove the key
        print(" * Remove the key %s, as the token is inactive!" % client_comment)
        res = salt_client.cmd(minion,
                              "ssh.rm_auth_key", 
                              kwarg=params)
        print " * %s" % res
    else:
        # add the key  
        print(" * Adding key %s to user %s" % (client_comment,
                                            user))
        params["comment"] =  client_comment
        params["enc"] = client_enc
        res = salt_client.cmd(minion,
                              "ssh.set_auth_key", 
                              kwarg=params)
        print " * %s" % res

##### main

def main():
    args = create_arguments()

    if not args.password:
        password = getpass.getpass(prompt="Please enter password for"
                                   " '%s':" % args.admin)
    else:
        password = args.password

    protocol = args.url.split(":", 1)[0]
    host = args.url.split(":", 1)[1].strip("/")
    # Create the privacyideaclient instance
    client = privacyideaclient(protocol,
                               host,
                               admin=args.admin,
                               adminpw=password,
                               adminrealm=args.adminrealm)
            
    params = {"application": "ssh"}   
    ret = client.connect("/machine/gettokenapps", {}, params)
    result = ret.get("result")
    
    if result.get("status"):
        total = result.get("value").get("total")
        print "Found %i keys." % total 
        machines = result.get("value").get("machines")
        
        '''
        TODO: we should work per machine. 
        1st delete the keys of the machine and 
        2nd then push new new ones..
        '''
        if args.purgekeys:
            machine_user_dict = get_machine_user_dict(machines)
            # print machine_user_dict
            for minion, users in machine_user_dict.iteritems():
                ssh_purge_keys(minion, users)
                
        for machine_id in machines.keys():
            set_ssh(machines[machine_id],
                    purgekeys=args.purgekeys)
    else:
        print "An error occurred!"
        print result


if __name__ == '__main__':
    if DEBUG:
        main()
    else:
        try:
            main()
        except Exception as ex:
            print "Error: %s" % ex
            sys.exit(5)
