Interactive plum run

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)

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()