#!/usr/bin/env python
import paramiko
import logging

###################################################################
## ATFSsh.Ssh Class contains variables and methods for 
## Connecting to remote machine using SSH
##
## Variables:
##      connections: List of Dictionaries.
##      Each Dictionary contains {host, user, conn} information.
##              host : Server IP address
##              user : User on host
##              conn : SSH Conn Object to user@host
###################################################################
class Ssh():
    

    def __init__(self):
        self.connections = []

    def connect(self, host, user):
        """
        Objective:
            SSH to Server "host" as User "user"

        Parameter:
            host: Server IP Address 
            user: Login Username

        Return:
            Success: 0
            Failure: 1
        """

        logger = logging.getLogger('ATF_LOG')
        connection = paramiko.SSHClient()
        connection.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        try:
            connection.connect(host, username=user)

        except paramiko.BadHostKeyException as result:
            logger.exception(
                "BadHostKeyException: Unable to Connect to Server: '" + host + 
                "' as User: '" + user + "'")
            return 1

        except paramiko.AuthenticationException:
            logger.exception("AuthenticationException: Unable to Authenticate "
                            + user + "@" + host)
            return 1

        except paramiko.SSHException:
            logger.exception("SSHException: Unknown server " + host)
            return 1

        else:
            self.connections.append({'host':host, 'user': user, 'conn':
                                    connection})
            logger.info("Successfully Able to SSH to: " + user + "@" + host)

        return 0 

    def close(self, host, user):
        """
        Objective:
            Close SSH Connections for User "username" 
            on the Server "server"

        Parameters:
            server: Server IP Address
            username: User on the Server

        Return:
            Success: 0
            Failure: 1
        """

        logger = logging.getLogger('ATF_LOG')
        index = 0
        
        for conn in self.connections:
            if conn['host'] == host and conn['user'] == user:
                conn['conn'].close()
                logger.info("Closing SSH Connection: " + conn['user'] + "@" +
                            conn['host'])         
                self.connection[index:(index + 1)] = []
                index = index - 1

            index = index + 1
    
        return 0

    def closeall(self):
        """
        Objective:
            Close All Existing SSH Connections

        Parameters:
            None. 

        Return:
            Success: 0
            Failure: 1
        """

        logger = logging.getLogger('ATF_LOG')
        
        for conn in self.connections:
            conn['conn'].close()
            logger.info("Closing SSH Connection: " + conn['user'] + "@" +
                        conn['host'])

        self.connections[:] = []

        return 0
    
    def getconnection(self, host, user):
        """
        Objective:
            Return SSH connection object for username@server

        Parameters:
            host: Server IP Address
            user: User on the Server

        Return:
            Success:SSH connection Object If SSH Connection Object exists 
                    for username@server
            Failure: 1 (If SSH Connection doesn't exist)
        """

        for conn in self.connections:
            if conn['host'] == host and conn['user'] == user:
                return conn['conn']

        else:
            return 1

    def executecommand(self, command, host, user):
        """
        Objective:
            Execute Command "comamnd" on user@host

        Parameters:
            host: IP address of Server/Client
            user: username of user on Host 'host'
            command: command to execute

        Return:
            Success: 0
            Failure: 1
        """

        logger = logging.getLogger('ATF_LOG')    
        output = []
        conn = self.getconnection(host, user)
    
        if conn == 1:
            logger.error("SSH Connection Not Established to:" + user + "@" +
                        host)
            output = [1, 0, 0, 0]
            return output
                        
        try:
            stdin, stdout, stderr = conn.exec_command(command)

        except paramiko.SSHException:
            logger.exception("Unable to Execute Command: " + command)
            output = [1, 0, 0, 0]
            return output
        
        else:
            output = [0, stdin, stdout, stderr]
            return output