#!/usr/bin/env python
#
# plessh will run the given shell command on all Planetlab Europe
# nodes.
#

import sys
import os
import xmlrpclib
from threading import Thread
from getpass import getpass

PLC_HOST = "www.planet-lab.eu"
SLICE_NAME = ""
KEY_FILE = "~/.ssh/id_rsa"
BASE_CMD = 'ssh -q %(options)s -i %(key)s -l %(slice_name)s %(host)s "%(cmd)s" 2> /dev/null'
THREAD_COUNT = 10
SSH_OPTIONS = {
    "BatchMode": "yes",
    "StrictHostKeyChecking": "no",
    "ConnectTimeout": 15,
    "UserKnownHostsFile": "/dev/null",
    "CheckHostIP": "no"
    }

class API:
    def __init__(self, hostname, username, password):
        self.host = hostname
        self.api = xmlrpclib.ServerProxy("https://%s/PLCAPI/" % self.host, allow_none=True)
        self.auth = {
            "Username": username,
            "AuthString": password,
            "AuthMethod": "password"
            }
        # this will raise an exception if not successful
        self.api.AuthCheck(self.auth)

    def wrap(self, function):
        def wrapper(*args):
            args = (self.auth, ) + args
            return function(*args)
        return wrapper

    def __getattr__(self, attr):
        return self.wrap(getattr(self.api, attr))


def getAPI(hostname, username=None, password=None):
    if not username: username=raw_input("Please enter your PLC username: ")
    if not password: password=getpass("Please enter your PLC password: ")
    while True:
        try:
            return API(hostname=hostname, username=username, password=password)
        except xmlrpclib.Fault,e:
            print e
            print "Please try again."
            username=raw_input("Please enter your PLC username: ")
            password=getpass("Please enter your PLC password: ")
            continue


def change_key_comment(output, hostname):
    fields = output.split()
    new_fields = []

    in_key = False
    in_comment = False
    for f in fields:
        f = f.strip()
        if f.startswith("ssh-"):
            in_key = True
        elif in_key:
            in_key = False
            in_comment = True
        elif in_comment:
            f = hostname
        new_fields.append(f)

    return " ".join(new_fields)
    


class Command(Thread):
    def __init__(self, cmd, hosts=[]):
        Thread.__init__(self)
        self.hosts = hosts
        self.cmd = cmd
        self.results = {}

    def runCmd(self, host):
        cmd = BASE_CMD % {'cmd': self.cmd,
                          'host': host,
                          'key': KEY_FILE,
                          'slice_name': SLICE_NAME,
                          'options':  " ".join(["-o %s=%s" % (o,v) for o,v in SSH_OPTIONS.items()])}
        p = os.popen(cmd)
        output = p.read().strip()
        if output:
            output = output.split('\n')[0] # only get the first line
# keep the comment for the keys
#            output = change_key_comment(output, host)
            return output
        return None

    def run(self):
        for host in self.hosts:
            self.results[host] = self.runCmd(host)


def distributeHosts(num_hosts):
    dist_list = []
    host_per_thread = num_hosts / THREAD_COUNT
    rest = num_hosts % THREAD_COUNT
    for i in range(THREAD_COUNT):
        c = host_per_thread
        if rest:
            c += 1
            rest -= 1
        dist_list.append(c)
    return dist_list


def main(plchost, slice_name, print_report=lambda:None):
    ple = getAPI(plchost)
    slice_nodes = []
    try:
        node_ids = ple.GetSlices(slice_name)[0]['node_ids']
        slice_nodes = [n['hostname'] for n in ple.GetNodes(node_ids)]
    except IndexError:
        print "Can not get nodes for slice: %s" % slice_name
        sys.exit(1)

    dist_list = distributeHosts(len(slice_nodes))
    thread_list = []
    index = 0
    cmd = "cat /home/%s/.ssh/id_rsa.pub" % slice_name
    print "Please wait gathering public keys from %d PlanetLab nodes. This may take a while..." % len(slice_nodes)
    for i in range(THREAD_COUNT):
        current = Command(cmd, slice_nodes[index:index+dist_list[i]])
        index += dist_list[i]
        thread_list.append(current)
        current.start()
    for i in thread_list:
        i.join()

    results = {}
    for i in thread_list:
        results.update(i.results)

    print_report(results)

if __name__ == "__main__":
    def save_report(results, fname="slice_keys"):
        ok = 0
        f = open(fname, "w")
        for node in results.keys():
            if results[node]:
                f.write("%s\n" % results[node])
                ok += 1
        f.close()
        print "Could get keys from %d of %d nodes" % (ok, len(results))
        print "Please see %s file." % fname

    def get_option(msg, default):
        ret = raw_input("%s [%s] : " % (msg, default))
        if not ret:
            ret = default
        return ret
            
    PLC_HOST=get_option("Please enter your PLC host", PLC_HOST)
    SLICE_NAME=get_option("Please enter your slice name", SLICE_NAME)
    KEY_FILE=get_option("Please enter the path of your key", KEY_FILE)
    main(PLC_HOST, SLICE_NAME, print_report=save_report)