#! /usr/bin/env python
# >><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><< #
# ssh - simple ssh2 wrapper
# Author: Aaron Ghent
# >><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><< #
#
# import ssh
# client = ssh.Connection(host='example.com', username='root', password='default')
# client.execute('ls')
# client.close()
#
# >><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><< #

import os
import tempfile
import paramiko

class Connection(object):
    def __init__(self,
            host,
            username=None,
            private_key=None,
            password=None,
            port=22,
            logfile=None,
        ):

        if not username:
            username = os.environ['LOGNAME']

        # Log to a temporary file.
        if logfile:
            paramiko.util.log_to_file(logfile)

        self._ssh = paramiko.SSHClient()
        self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self._ssh.load_system_host_keys()
        self._ssh.connect(hostname=host, username=username, password=password)

        # Authenticate the transport.
        if password:
            # Using Password.
            self._ssh.connect(hostname=host, username=username, password=password)
        else:
            # Use Private Key.
            if not private_key:
                    # Try to use default key.
                    if os.path.exists(os.path.expanduser('~/.ssh/id_rsa')):
                            private_key = '~/.ssh/id_rsa'
                    elif os.path.exists(os.path.expanduser('~/.ssh/id_dsa')):
                            private_key = '~/.ssh/id_dsa'
                    else:
                            raise TypeError, "You have not specified a password or key."

            private_key_file = os.path.expanduser(private_key)
            rsa_key = paramiko.RSAKey.from_private_key_file(private_key_file)
            self._ssh.connect(hostname=host, username=username, pkey=rsa_key)

    def get(self, remotepath, localpath = None):
        """
        Copies a file between the remote host and the local host.
        """

        if not localpath:
                localpath = os.path.split(remotepath)[1]

        sftp = self._ssh.open_sftp()
        sftp.get(remotepath, localpath)

    def put(self, localpath, remotepath = None):
        """
        Copies a file between the local host and the remote host.
        """
        if not remotepath:
                remotepath = os.path.split(localpath)[1]

        sftp = self._ssh.open_sftp()
        sftp.put(localpath, remotepath)

    def execute(self, command):
        """
        Execute the given commands on a remote machine.
        """
        try:
                std_in, std_out, std_err = self._ssh.exec_command(command)
        except:
                return '[Error] execute: command failed'

        if std_out.readlines():
                return std_out.readlines()
        elif std_err.readlines():
                return std_err.readlines()
        else:
                return

    def close(self):
        """
        Closes the connection and cleans up.
        """
        self._ssh.close()

    def __del__(self):
        """
        Attempt to clean up if not explicitly closed.
        """
        self.close()

# >><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><<>><< #