First draft of econtext/sudo.py.

docs/history.rst

@@ -53,6 +53,7 @@ contribution to a solution, I hope you find it useful.
 Future
 ######
 
+* Move connect() logic to dispatcher thread.
 * Connect back using TCP and SSL.
 * Python 3 support.
 * Windows support via psexec or similar.

econtext/master.py

@@ -11,9 +11,11 @@ import logging
 import os
 import pkgutil
 import re
+import select
 import socket
 import sys
 import textwrap
+import time
 import types
 import zlib
 
@@ -68,16 +70,29 @@ def create_child(*args):
     return pid, os.dup(parentfp.fileno())
 
 
-def discard_until(fd, s):
+def read_with_deadline(fd, size, deadline):
+    timeout = deadline - time.time()
+    if timeout > 0:
+        rfds, _, _ = select.select([fd], [], [], timeout)
+        if rfds:
+            return os.read(fd, size)
+
+    raise econtext.core.TimeoutError('read timed out')
+
+def iter_read(fd, deadline):
     buf = ''
-    while not buf.endswith(s):
-        print 'hi'
+    while True:
         s = os.read(fd, 4096)
-        print ['got', s]
         if not s:
-            raise econtext.core.StreamError('Expected %r, received %r', s, buf)
-
+            raise econtext.core.StreamError('EOF on stream; received %r', buf)
         buf += s
+        yield buf
+
+
+def discard_until(fd, s, deadline):
+    for buf in iter_read(fd, deadline):
+        if buf.endswith(s):
+            return
 
 
 class LogForwarder(object):
@@ -282,14 +297,28 @@ class Stream(econtext.core.Stream):
         LOG.debug('%r.connect(): child process stdin/stdout=%r',
                   self, self.receive_side.fd)
 
-        discard_until(self.receive_side.fd, 'EC0\n')
+        self._connect_bootstrap()
+
+    def _ec0_received(self):
+        LOG.debug('%r._ec0_received()', self)
         econtext.core.write_all(self.transmit_side.fd, self.get_preamble())
-        discard_until(self.receive_side.fd, 'EC1\n')
+        discard_until(self.receive_side.fd, 'EC1\n', time.time() + 10.0)
+
+    def _connect_bootstrap(self):
+        discard_until(self.receive_side.fd, 'EC0\n', time.time() + 10.0)
+        return self._ec0_received()
 
 
 class Broker(econtext.core.Broker):
     shutdown_timeout = 5.0
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, e_type, e_val, tb):
+        self.shutdown()
+        self.join()
+
 
 class Context(econtext.core.Context):
     def __init__(self, *args, **kwargs):

econtext/sudo.py

@@ -0,0 +1,123 @@
+
+import logging
+import os
+import pty
+import termios
+import time
+
+import econtext.master
+
+
+LOG = logging.getLogger(__name__)
+PASSWORD_PROMPT = 'password'
+
+
+def flags(names):
+    """Return the result of ORing a set of (space separated) :py:mod:`termios`
+    module constants together."""
+    return sum(getattr(termios, name) for name in names.split())
+
+
+def cfmakeraw((iflag, oflag, cflag, lflag, ispeed, ospeed, cc)):
+    """Given a list returned by :py:func:`termios.tcgetattr`, return a list
+    that has been modified in the same manner as the `cfmakeraw()` C library
+    function."""
+    iflag &= ~flags('IGNBRK BRKINT PARMRK ISTRIP INLCR IGNCR ICRNL IXON')
+    oflag &= ~flags('OPOST IXOFF')
+    lflag &= ~flags('ECHO ECHOE ECHONL ICANON ISIG IEXTEN')
+    cflag &= ~flags('CSIZE PARENB')
+    cflag |= flags('CS8')
+
+    iflag = 0
+    oflag = 0
+    lflag = 0
+    return [iflag, oflag, cflag, lflag, ispeed, ospeed, cc]
+
+
+def disable_echo(fd):
+    old = termios.tcgetattr(fd)
+    new = cfmakeraw(old)
+    #new = old[:]
+    #new[3] &= ~flags('ECHO')
+    tcsetattr_flags = termios.TCSAFLUSH
+    if hasattr(termios, 'TCSASOFT'):
+        tcsetattr_flags |= termios.TCSASOFT
+    termios.tcsetattr(fd, tcsetattr_flags, new)
+
+
+def tty_create_child(*args):
+    master_fd, slave_fd = os.openpty()
+    import econtext.core
+    #econtext.core.set_nonblocking(master_fd)
+    disable_echo(master_fd)
+    disable_echo(slave_fd)
+
+    pid = os.fork()
+    if not pid:
+        os.dup2(slave_fd, 0)
+        os.dup2(slave_fd, 1)
+        os.dup2(slave_fd, 2)
+        os.close(slave_fd)
+        os.close(master_fd)
+        #os.setsid()
+        os.close(os.open(os.ttyname(1), os.O_RDWR))
+        os.execvp(args[0], args)
+        raise SystemExit
+
+    os.close(slave_fd)
+    LOG.debug('tty_create_child() child %d fd %d, parent %d, args %r',
+              pid, master_fd, os.getpid(), args)
+    return pid, master_fd
+
+
+class StateMachine(object):
+    S_PW_PROMPT, S_SHELL, S_CONNECTED = range(3)
+    state = S_PW_PROMPT
+
+
+class Stream(econtext.master.Stream):
+    create_child = staticmethod(tty_create_child)
+    sudo_path = 'sudo'
+    password = None
+
+    def get_boot_command(self):
+        bits = [self.sudo_path, '-S', '-u', self._context.username]
+        return bits + super(Stream, self).get_boot_command()
+
+    def _connect_bootstrap(self):
+        password_sent = False
+        it = econtext.master.iter_read(self.receive_side.fd,
+                                       time.time() + 10.0)
+
+        for buf in it:
+            if buf.endswith('EC0\n'):
+                return self._ec0_received()
+            elif (not password_sent) and 'password' in buf.lower():
+                if self.password is None:
+                    raise econtext.core.StreamError('password required')
+                LOG.debug('sending password')
+                os.write(self.transmit_side.fd, self.password + '\n')
+                password_sent = True
+        else:
+            raise econtext.core.StreamError('bootstrap failed')
+
+
+def connect(broker, username=None, sudo_path=None, python_path=None, password=None):
+    """Get the named sudo context, creating it if it does not exist."""
+    if username is None:
+        username = 'root'
+
+    context = econtext.master.Context(
+        broker=broker,
+        name='sudo:' + username,
+        username=username)
+
+    context.stream = Stream(context)
+    if sudo_path:
+        context.stream.sudo_path = sudo_path
+    if password:
+        context.stream.password = password
+    if python_path:
+        context.stream.python_path = python_path
+    context.stream.connect()
+    return broker.register(context)