connection.py 3.18 KB
Newer Older
Cedric Roux's avatar
Cedric Roux committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
import os, subprocess, time, fcntl, termios, tty, signal, thread

from utils import log

class connection:
    def __init__(self, description, host, user, password):
        self.description = description
        self.host = host
        self.user = user
        self.password = password
        self.sendlock = thread.allocate_lock()

        try:
            (pid, fd) = os.forkpty()
        except BaseException, e:
            log("ERROR: forkpty for '" + description + "': " + e)
            (pid, fd) = (-1, -1)

        if pid == -1:
            log("ERROR: creating connection for '" + description + "'")
            os._exit(1)

        # child process, run ssh
        if pid == 0:
            try:
                os.execvp('sshpass', ['sshpass', '-p', password,
                          'ssh', user + '@' + host])
            except BaseException, e:
                log("ERROR: execvp for '" + description + "': " + e)
            log("ERROR: execvp failed for '" + description + "'")
            os._exit(1)

        # parent process
        # make the TTY raw to avoid getting input printed and no ^M
        try:
            tty.setraw(fd, termios.TCSANOW)
        except BaseException, e:
            log("ERROR: failed configuring TTY: " + str(e))
            os._exit(1)

#        try:
#            fcntl.fcntl(fd, fcntl.F_SETFL,
#                        fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK)
#        except:
#            log("ERROR: fcntl failed for '" + description + "'")
#            os._exit(1)

        self.pid = pid
        self.fd = fd
        self.active = True
        self.retcode = -1

    def send(self, string):
        if self.active == False:
            log("ERROR: send: child is dead for '" + self.description + "'")
            return -1

        try:
            (pid, out) = os.waitpid(self.pid, os.WNOHANG)
        except BaseException, e:
            log("ERROR: send: waitpid failed for '" + self.description +
                "': " + str(e))
            (pid, out) = (self.pid, 1)
        if pid != 0:
            log("ERROR: send: child process dead for '" +
                self.description + "'")
            try:
                os.close(self.fd)
            except BaseException, e:
                #we don't care about errors at this point
                pass
            self.active = False
            self.retcode = out
            return -1

        self.sendlock.acquire()

        length = len(string)
        while length != 0:
            try:
                ret = os.write(self.fd, string)
            except BaseException, e:
                log("ERROR: send fails for '" + self.description + "': " +
                    str(e))
                os._exit(1)

            if ret == 0:
                log("ERROR: send: write returns 0 for '" +
                    self.description + "'")
                os._exit(1)

            length = length - ret
            string = string[ret:]

        self.sendlock.release()

        return 0

    def kill(self, signal=signal.SIGTERM):
        log("INFO: kill connection '" + self.description + "'")
        try:
            os.kill(self.pid, signal)
        except BaseException, e:
            log("ERROR: connection.kill: " + str(e))