# Copyright (C) 2008  Justin Cook <jhcook_AT_gmail.com>
# Some code (C) 2003-2007 Robey Pointer <robey@lag.net>
#
# This 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 2.1 of the License, or (at your option)
# any later version.
#
# It is distrubuted 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.
#
# If you do not have a copy of the GNU Lesser General Public License,
# you can write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.

import paramiko
import traceback
import binascii
import socket
import logging
import time
import os

class ConnectionError(Exception):
    def __init__(self, value):
        self.__value = value
    def __str__(self):
        return self.__value

class Connection:
    """
    """
    def __init__(self, hostname, port):
        self.__hostname = hostname
        self.__port = port

    def connect(self, username, password):
        self.__username = username
        self.__password = password
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((self.__hostname, self.__port))
        except Exception, e:
            raise ConnectionError(str(e))
        try:
            self.__t = paramiko.Transport(sock)
            try:
                self.__t.start_client()
            except paramiko.SSHException:
                raise ConnectionError('SSH negotiation failed for %s' % self.__hostname)
            try:
                keys = paramiko.util.load_host_keys(os.path.expanduser('~/.ssh/known_hosts'))
            except IOError:
                try:
                    keys = paramiko.util.load_host_keys(os.path.expanduser('~/ssh/known_hosts'))
                except IOError:
                    keys = {}
            # check server's host key -- this is important.
            key = self.__t.get_remote_server_key()
            if not keys.has_key(self.__hostname):
                raise ConnectionError ('unknown host key for %s' % self.__hostname)
            elif not keys[self.__hostname].has_key(key.get_name()):
                raise ConnectionError('unknown host key on %s' % self.__hostname)
            elif keys[self.__hostname][key.get_name()] != key:
                raise ConnectionError('%s key has changed' % self.__hostname)
            else:
                logging.debug('%s key OK' % self.__hostname)
            self.__agent_auth()
            if not self.__t.is_authenticated():
                self.__t.close()
                raise ConnectionError('authentiation failed on %s' % self.__hostname)
        except Exception, e:
            raise ConnectionError('caught exception: ' + str(e.__class__) + ': ' + str(e))
            traceback.print_exc()
            try:
                self.__t.close()
            except:
                pass
            return False
        return True

    def __agent_auth(self):
        """
        Attempt to authenticate to the given transport using any of the private
        keys available from an SSH agent.
        """
        agent = paramiko.Agent()
        agent_keys = agent.get_keys()
        if len(agent_keys) == 0:
            return
        for key in agent_keys:
            logging.debug('Trying ssh-agent key %s' % binascii.hexlify(key.get_fingerprint()))
            try:
                self.__t.auth_publickey(self.__username, key)
                logging.debug('Key accepted on %s' % self.__hostname)
                return
            except paramiko.SSHException:
                raise ConnectionError('Key not accepted on %s' % self.__hostname)

    def sudo_command(self, command, **kw):
        """
        We should have a valid connection. If not raise an exception.
        Attempt to establish a transport over the channel otherwise.
        """
        # Check if we have an authenticated session:
        if not self.__t.is_authenticated():
            raise ConnectionError('transport not authenticated on %s' % self.__hostname)
        # Check arguments:
        if kw.has_key('maxtry'):
            maxtry = kw.get('maxtry')
        else:
            maxtry = 3
        if kw.has_key('args'):
            command_string = 'sudo -S -p [%%u@sudo] %s %s' % (kw.get('args'), command)
        else:
            command_string = 'sudo -S -p [%%u@sudo] %s' % command
        # Open the session:
        chan = self.__t.open_session()
        # Execute the command:
        chan.exec_command(command_string)
        # Wait until we have recv_ready. If not recv_ready, we check for recv_stderr_ready.
        # In case it is available, we read and check to see if it is the password prompt.
        # If so we pass the password.
        passcount = 0
        while not chan.recv_ready():
            if passcount == maxtry:
                chan.close()
                raise ConnectionError('bad password on %s' % self.__hostname)
            elif chan.recv_stderr_ready():
                stderr = chan.recv_stderr(1024)
                if stderr.find('%s@sudo' % self.__username) > 0:
                    chan.send('%s\r\n' % self.__password)
                    passcount += 1
                elif stderr.find('not in the sudoers') > 0:
                    raise ConnectionError('%s not in the sudoers file on %s' %
                                          (self.__username, self.__hostname))
                elif stderr.find('not allowed to execute') > 0:
                    raise ConnectionError('%s is not allowed to execute %s on %s' %
                                          (self.__username, command, self.__hostname))
                else:
                    chan.close()
                    logging.critical(stderr)
                    raise ConnectionError('unknown sudo error on %s' % self.__hostname)
            else:
                time.sleep(.1)
        while chan.recv_ready():
            data = chan.recv(1024)
            logging.debug('Received from %s: \r\n%s' % (self.__hostname, data))
            time.sleep(.1)
        return chan.recv_exit_status()