#!/usr/bin/env python

import base64
import paramiko
from argparse import ArgumentParser
from ConfigParser import ConfigParser
from Crypto.Cipher import AES

class NodeMonitor(object):
    """
    Class for running commands on a remote host using SSH.
    """

    def __init__(self, config, keyfile):
        """
        Constructor. Sets the configuration object and creates the cipher for
        password/user encryption.
        :param config: ConfigParser
        :param keyfile: path to a file which contains the encryption key
        """
        self.config = config
        with open(keyfile) as KEY:
            key = KEY.read()
        self.cipher = AES.new(key, AES.MODE_ECB)

    def monitor(self, nodefilter, commandfilter, customcommand):
        """
        Runs the configured or custom commands on the given or on all configured
        nodes.
        :param nodefilter: list of nodes [str]
        :param commandfilter: list of command-keys [str]
        :param customcommand: custom command str
        """
        for node in self.getNodes(nodefilter):
            try:
                connection = self.connect(
                    node['host'],
                    node['user'],
                    node['pwd']
                )
                for cmd in self.getCommands(node, commandfilter, customcommand):
                    print('\n{}: {}'.format(node['host'].upper(), cmd))
                    try:
                        self.execute(connection, cmd)
                    except:
                        print('error. command could not be executed.')
            finally:
                connection.close()

    def connect(self, host, username, password):
        """
        Establishes a SSH connection.
        :param host: hostname str
        :param username: username str
        :param password: password str
        :return paramiko ssh connection
        """
        username = self.decrypt(username)
        password = self.decrypt(password)
        connection = paramiko.SSHClient()
        connection.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        connection.connect(host, username=username, password=password)

        return connection

    def execute(self, connection, cmd):
        """
        Executes the given command using the provides paramiko ssh connection.
        :param connection: paramiko ssh connection
        :param cmd: command str
        """
        stdin, stdout, stderr = connection.exec_command(cmd)
        stdin.close()

        try:
            for line in stdout.readlines():
                print(line)
            for line in stderr.readlines():
                print(line)
        finally:
            stdout.close()
            stderr.close()

    def getNodes(self, nodefilter=[]):
        """
        Returns a list of node information (dict).
        :param nodefilter: list of nodes to include [str]
        :return list of node information [{}]
        """
        nodes = []
        for section in self.config.sections():
            if section.startswith('node'):
                node = {key: value for key, value in self.config.items(section)}
                if not nodefilter or node['host'] in nodefilter:
                    nodes.append(node)

        return nodes

    def getCommands(self, node, cmdfilter=None, customcommand=None):
        """
        Returns a list of commands for the given node.
        :param node: node on which the commands should be run
        :param cmdfilter: list of command-keys [str]
        :param customcommand: custom command str
        :return list of commands [str]
        """
        cmds = []
        items = node.items()
        if cmdfilter:
            for filtr in cmdfilter:
                for key, value in items:
                    if key.startswith('cmd.{}'.format(filtr)):
                        cmds.append(value)
        elif not customcommand:
            cmds = [x[1] for x in filter(lambda x: x[0].startswith('cmd.'), items)]

        if customcommand:
            cmds.append(customcommand)

        return cmds

    def listConfig(self):
        """
        Prints a list of available commands for each node.
        """
        for node in self.getNodes():
            print('\n{}'.format(node['host'].upper()))
            for key, value in node.items():
                if key.startswith('cmd.'):
                    print('{}: {}'.format(key.replace('cmd.', ''), value))

    def encrypt(self, password):
        """
        Encrypts the given password.
        :param password: str
        """
        print(base64.b64encode(self.cipher.encrypt(password.rjust(32))))

    def decrypt(self, password):
        """
        Decrypts the given password.
        :param password: str
        :return decrypted password str
        """
        return self.cipher.decrypt(base64.b64decode(password)).strip()


if __name__ == '__main__':

    argParser = ArgumentParser(description='NodeMonitor - run commands on \
                                            remote hosts')
    argParser.add_argument('key', type=str, metavar='KEY',
                           help='keyfile for encrypting/decrypting passwords')
    argParser.add_argument('-c', '--cmds', nargs='+', type=str, metavar='CMD',
                           help='run given configured commands')
    argParser.add_argument('-C', '--custom', type=str, metavar='CMD',
                           help='run given custom command')
    argParser.add_argument('-e', '--encrypt', type=str, metavar='PWD',
                           help='encrypt the given password')
    argParser.add_argument('-n', '--nodes', nargs='+', type=str, metavar='NODE',
                           help='run commands on given nodes only')
    argParser.add_argument('-l', '--list', action='store_true',
                           help='list available nodes and commands')
    args = vars(argParser.parse_args())

    config = ConfigParser()
    config.read('nodemonitor.cfg')

    nodeMonitor = NodeMonitor(config, args['key'])

    if args['list']:
        nodeMonitor.listConfig()
    elif args['encrypt']:
        nodeMonitor.encrypt(args['encrypt'])
    else:
        nodeMonitor.monitor(args['nodes'], args['cmds'], args['custom'])