From abfb3b800f7bcd5380189a34db8629270d4708e5 Mon Sep 17 00:00:00 2001 From: Ben Firshman Date: Fri, 20 Dec 2013 15:03:01 +0000 Subject: [PATCH] Interactive plum run --- plum/cli/main.py | 48 +++++++++++++-- plum/cli/socketclient.py | 129 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 171 insertions(+), 6 deletions(-) create mode 100644 plum/cli/socketclient.py diff --git a/plum/cli/main.py b/plum/cli/main.py index 80646e339..b3e4f11e6 100644 --- a/plum/cli/main.py +++ b/plum/cli/main.py @@ -12,6 +12,7 @@ from .log_printer import LogPrinter from docker.client import APIError from .errors import UserError from .docopt_command import NoSuchCommand +from .socketclient import SocketClient log = logging.getLogger(__name__) @@ -122,18 +123,22 @@ class TopLevelCommand(Command): raise UserError("No such service: %s" % options['SERVICE']) container_options = { 'command': [options['COMMAND']] + options['ARGS'], + 'tty': not options['-d'], + 'stdin_open': not options['-d'], } container = service.create_container(one_off=True, **container_options) if options['-d']: service.start_container(container, ports=None) print container.name else: - stream = container.logs(stream=True) - service.start_container(container, ports=None) - for data in stream: - if data is None: - break - print data + with self._attach_to_container( + container.id, + interactive=True, + logs=True, + raw=True + ) as c: + service.start_container(container, ports=None) + c.run() def start(self, options): """ @@ -185,6 +190,37 @@ class TopLevelCommand(Command): print "Attaching to", list_containers(containers) LogPrinter(containers, attach_params={'logs': True}).run() + def _attach_to_container(self, container_id, interactive, logs=False, stream=True, raw=False): + stdio = self.client.attach_socket( + container_id, + params={ + 'stdin': 1 if interactive else 0, + 'stdout': 1, + 'stderr': 0, + 'logs': 1 if logs else 0, + 'stream': 1 if stream else 0 + }, + ws=True, + ) + + stderr = self.client.attach_socket( + container_id, + params={ + 'stdin': 0, + 'stdout': 0, + 'stderr': 1, + 'logs': 1 if logs else 0, + 'stream': 1 if stream else 0 + }, + ws=True, + ) + + return SocketClient( + socket_in=stdio, + socket_out=stdio, + socket_err=stderr, + raw=raw, + ) def list_containers(containers): return ", ".join(c.name for c in containers) diff --git a/plum/cli/socketclient.py b/plum/cli/socketclient.py new file mode 100644 index 000000000..90ed8b587 --- /dev/null +++ b/plum/cli/socketclient.py @@ -0,0 +1,129 @@ +# Adapted from https://github.com/benthor/remotty/blob/master/socketclient.py + +from select import select +import sys +import tty +import fcntl +import os +import termios +import threading +import errno + +import logging +log = logging.getLogger(__name__) + + +class SocketClient: + def __init__(self, + socket_in=None, + socket_out=None, + socket_err=None, + raw=True, + ): + self.socket_in = socket_in + self.socket_out = socket_out + self.socket_err = socket_err + self.raw = raw + + self.stdin_fileno = sys.stdin.fileno() + + def __enter__(self): + self.create() + return self + + def __exit__(self, type, value, trace): + self.destroy() + + def create(self): + if os.isatty(sys.stdin.fileno()): + self.settings = termios.tcgetattr(sys.stdin.fileno()) + else: + self.settings = None + + if self.socket_in is not None: + self.set_blocking(sys.stdin, False) + self.set_blocking(sys.stdout, True) + self.set_blocking(sys.stderr, True) + + if self.raw: + tty.setraw(sys.stdin.fileno()) + + def set_blocking(self, file, blocking): + fd = file.fileno() + flags = fcntl.fcntl(fd, fcntl.F_GETFL) + flags = (flags & ~os.O_NONBLOCK) if blocking else (flags | os.O_NONBLOCK) + fcntl.fcntl(fd, fcntl.F_SETFL, flags) + + def run(self): + if self.socket_in is not None: + self.start_background_thread(target=self.send_ws, args=(self.socket_in, sys.stdin)) + + recv_threads = [] + + if self.socket_out is not None: + recv_threads.append(self.start_background_thread(target=self.recv_ws, args=(self.socket_out, sys.stdout))) + + if self.socket_err is not None: + recv_threads.append(self.start_background_thread(target=self.recv_ws, args=(self.socket_err, sys.stderr))) + + for t in recv_threads: + t.join() + + def start_background_thread(self, **kwargs): + thread = threading.Thread(**kwargs) + thread.daemon = True + thread.start() + return thread + + def recv_ws(self, socket, stream): + try: + while True: + chunk = socket.recv() + + if chunk: + stream.write(chunk) + stream.flush() + else: + break + except Exception, e: + log.debug(e) + + def send_ws(self, socket, stream): + while True: + r, w, e = select([stream.fileno()], [], []) + + if r: + chunk = stream.read(1) + + if chunk == '': + socket.send_close() + break + else: + try: + socket.send(chunk) + except Exception, e: + if hasattr(e, 'errno') and e.errno == errno.EPIPE: + break + else: + raise e + + def destroy(self): + if self.settings is not None: + termios.tcsetattr(self.stdin_fileno, termios.TCSADRAIN, self.settings) + + sys.stdout.flush() + +if __name__ == '__main__': + import websocket + + if len(sys.argv) != 2: + sys.stderr.write("Usage: python socketclient.py WEBSOCKET_URL\n") + exit(1) + + url = sys.argv[1] + socket = websocket.create_connection(url) + + print "connected\r" + + with SocketClient(socket, interactive=True) as client: + client.run()