Commit 8d24a8c6 authored by Anael Beutot's avatar Anael Beutot
Browse files

Simplified RShell class + handle child process kill in close method

parent c2c514a9
Loading
Loading
Loading
Loading
+19 −41
Original line number Diff line number Diff line
@@ -79,22 +79,16 @@ class FakePtySocket(object):

class RemoteShell(object):
    """Handles basic operations on remote shell."""
    def __init__(self, ev_loop, conn, exec_='/bin/bash'):
    def __init__(self, conn, exec_='/bin/bash'):
        """
        :param ev_loop: pyev loop instance
        :param conn: sjRPC connection
        :param exec_: binary path

        .. warning:: constructor must be called inside the thread runnig libev
            loop, this is necessary in order to create and start watcher and to
            close file descriptors
        :param exec_: path of binary to execute
        """
        # handles close when waiting process to end from an other thread
        self.watcher = ev_loop.async(self.close_cb)
        self.watcher.start()
        self.master, self.slave = pty.openpty()
        self.endpoint = FakePtySocket(self.master)
        self.proto = conn.create_tunnel(endpoint=self.endpoint)
        self.proto = conn.create_tunnel(endpoint=self.endpoint,
                                        on_close=self.on_tunnel_close)
        self.conn = conn
        try:
            self.process = subprocess.Popen(
@@ -116,23 +110,21 @@ class RemoteShell(object):
        return self.proto.label

    def resize(self, row, col, xpixel, ypixel):
        if self.master is not None:
            ioctl(self.master, termios.TIOCSWINSZ,
                  struct.pack('HHHH', row, col, xpixel, ypixel))

    def close(self):
        # async
        self.watcher.send()

    def terminate(self):
        self.process.terminate()
    def on_tunnel_close(self, *args):
        # *args is for callback arguments (ignored)
        self.proto, proto = None, self.proto  # prevents multiple calls to
                                              # close method
        self.close()

    def kill(self):
    def close(self):
        if self.process.poll() is None:
            # process is still alive
            self.process.kill()

    def close_cb(self, *args):
        if self.watcher is not None:
            self.watcher.stop()
            self.watcher = None
            self.process.wait()
        if self.master is not None:
            try:
                os.close(self.master)
@@ -152,14 +144,6 @@ class RemoteShell(object):
                logger.error('Error while trying to close RPC tunnel')
            self.proto = None

    def wait(self):
        return self.process.wait()

    def wait_n_close(self):
        ret_code = self.wait()
        self.close()
        return ret_code


class Handler(BasePlugin):
    """Handler for host role."""
@@ -248,13 +232,7 @@ class Handler(BasePlugin):
        shutil.rmtree(self.scripts_dir, ignore_errors=True)
        # kill all currently running shells
        for shell in self.shells.values():
            try:
                shell.kill()
            except OSError as exc:
                if exc.errno != errno.ESRCH:
                    raise
                # process does not exists
            shell.close_cb()
            shell.close()
        self.shells.clear()
        BasePlugin.stop(self)

@@ -318,7 +296,7 @@ class Handler(BasePlugin):
        """Create a shell tunnel and return the label of the created tunnel.

        """
        remote_shell = RemoteShell(self.main.evloop, conn, shell)
        remote_shell = RemoteShell(conn, shell)
        self.shells[remote_shell.label] = remote_shell
        return remote_shell.label