root/dual/trunk/dual.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2006 by KANEMOTO Shigeru

import os
import sys
import tty
import pty
import select
import struct
import fcntl
import termios
import errno
import signal
import socket
import getopt
import curses
import pwd

STDIN_FILENO = 0
STDOUT_FILENO = 1

ENV_FLAGNAME = 'DUALPY'

version = 'dual.py 2006/4/17 Copyright (c) 2006 by KANEMOTO Shigeru'
progname = sys.argv[0].split('/')[-1]


def getwinsz(fd):
  winsize = struct.pack('HHHH', 0, 0, 0, 0)
  winsize = fcntl.ioctl(fd, termios.TIOCGWINSZ, winsize)
  return struct.unpack('HHHH', winsize)
  # return (row, col, xpixel, ypixel)

def setwinsz(fd, list):
  # list = (row, col, xpixel, ypixel)
  winsize = struct.pack('HHHH', *list)
  fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)

xterm = 0

def settermwinsz(fd, list):
  # list = (row, col, hpixel, vpixel)
  global xterm
  if xterm:
    (row, col, xpixel, ypixel) = list
    os.write(fd, '\033[8;%d;%d;t\033[4;%d;%d;t' % (row, col, ypixel, xpixel))

#
# xterm specific escape sequences
#
# ESC [ 4 ; height ; width ; t
#   set screen size in pixels
# ESC [ 8 ; height ; width ; t
#   set screen size in characters
# ESC [ 14 ; height ; width ; t
#   request screen size in pixels
# ESC [ 18 ; height ; width ; t
#   request screen size in characters
#

owinsz = (0, 0, 0, 0)

def enter_mode():
  try:
    global owinsz, omode
    owinsz = getwinsz(STDOUT_FILENO)
    omode = tty.tcgetattr(STDIN_FILENO)
    tty.setraw(STDIN_FILENO)
  except tty.error, e:
    print >>sys.stderr, "%s: Can't get/set tty mode: %s" % (progname, e[1])
    sys.exit(os.EX_IOERR)

def reset_mode():
  global owinsz, omode
  tty.tcsetattr(STDIN_FILENO, tty.TCSAFLUSH, omode)
  if owinsz != (0, 0, 0, 0):
    settermwinsz(STDOUT_FILENO, owinsz)

def clearscreen():
  os.write(STDOUT_FILENO, curses.tigetstr('clear'))

def writeall(fd, data):
  while len(data) != 0:
    n = os.write(fd, data)
    data = buffer(data, n)

class Decoder:
  def __init__(self, sink, screen):
    self.sink_ = sink
    self.screen_ = screen
    self.first_ = 1
    self.s_ = ''

  def data(self, s):
    self.s_ += s
    while 1:
      if len(self.s_) < 3:
        return
      (length, cmd) = struct.unpack('!HB', self.s_[:3])
      if len(self.s_) < length:
        return

      s = self.s_[3:length]
      self.s_ = self.s_[length:]

      if cmd == 0:
        writeall(self.sink_, s)
      elif cmd == 1:
        (height, width) = struct.unpack('!HH', s[:4])
        (oheight, owidth, oxpixel, oypixel) = getwinsz(self.screen_)
        change = 0
        if self.first_:
          if height < oheight:
            change = 1
            oheight = height
          if width < owidth:
            change = 1
            owidth = width
          self.first_ = 0
        else:
          if height != oheight or width != owidth:
            change = 1
          oheight = height
          owidth = width
        if change:
          setwinsz(self.screen_, (oheight, owidth, oxpixel, oypixel))
          settermwinsz(self.screen_, (oheight, owidth, oxpixel, oypixel))

class Encoder:
  def __init__(self, dst):
    self.dst_ = dst

  def sink(self, s):
    self.dst_.send(struct.pack('!HB', len(s) + 3, 0) + s)

  def winsz(self, list):
    (height, width, xpixel, ypixel) = list
    self.dst_.send(struct.pack('!HBHH', 7, 1, height, width))


#
# Client
#

def client(dstport):
  try:
    print >>sys.stderr, '*** Connecting to port %d' % dstport
    link = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    link.connect(('127.0.0.1', dstport))
  except socket.error, e:
    print >>sys.stderr, "%s: Can't connect: %s" % (progname, e[1])
    sys.exit(os.EX_UNAVAILABLE)

  clearscreen()
  enter_mode()
  decoder = Decoder(STDOUT_FILENO, STDOUT_FILENO)
  encoder = Encoder(link)
  encoder.winsz(getwinsz(STDOUT_FILENO))

  def sigwinch(s, frame):
    encoder.winsz(getwinsz(STDOUT_FILENO))
  signal.signal(signal.SIGWINCH, sigwinch)

  try:
    try:
      while 1:
        try:
          (rfds, wfds, xfds) = select.select([STDIN_FILENO, link], [], [], 1)

          if STDIN_FILENO in rfds:
            encoder.sink(os.read(STDIN_FILENO, 1024))

          if link in rfds:
            data = link.recv(1024)
            if len(data) == 0:
              break
            decoder.data(data)

        except (select.error, OSError), e:
          if e[0] != errno.EINTR:
            raise e

    except OSError, e:
      if e[0] != errno.EIO:             # EIO when the link losts
        raise

  finally:
    reset_mode()
    print >>sys.stderr, '*** Closed the connection with port %d' % dstport


#
# Server
#

def server(argv):
  (pid, fd) = pty.fork()
  if pid == 0:
    # child
    try:
      os.environ[ENV_FLAGNAME] = ''
      os.execvp(argv[0], argv)
    except OSError, e:
      os.write(STDOUT_FILENO,
        '\0\0\0Can\'t exec "%s": %s\0\0\0' % (argv[0], e[1]))
      os._exit(os.EX_OSERR)

  data = os.read(fd, 1024)
  if data[0:3] == '\0\0\0':
    data = data[3:]
    data = data[:data.index('\0\0\0')]
    print >>sys.stderr, '%s: %s' % (progname, data)
    (pid, st) = os.waitpid(pid, 0)
    sys.exit(os.WIFEXITED(st) and os.WEXITSTATUS(st) or os.EX_OSERR)

  listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  listener.bind(('127.0.0.1', 0))       # localhost, any port
  print >>sys.stderr, '*** Waiting on port %d' % listener.getsockname()[1]
  listener.listen(1)
  (link, sa) = listener.accept()
  del listener
  link.setblocking(0)

  clearscreen()
  enter_mode()
  decoder = Decoder(fd, STDOUT_FILENO)
  setwinsz(fd, getwinsz(STDOUT_FILENO))
  encoder = Encoder(link)
  encoder.winsz(getwinsz(STDOUT_FILENO))
  encoder.sink(data)
  writeall(STDOUT_FILENO, data)

  def sigwinch(s, frame):
    winsz = getwinsz(STDOUT_FILENO)
    setwinsz(fd, winsz)
    encoder.winsz(winsz)
  signal.signal(signal.SIGWINCH, sigwinch)

  try:
    try:
      while 1:
        try:
          (rfds, x, x) = select.select([fd, STDIN_FILENO, link], [], [], 1)

          if fd in rfds:
            data = os.read(fd, 1024)
            writeall(STDOUT_FILENO, data)
            encoder.sink(data)

          if STDIN_FILENO in rfds:
            writeall(fd, os.read(STDIN_FILENO, 1024))

          if link in rfds:
            decoder.data(link.recv(1024))

        except (select.error, OSError), e:
          if e[0] != errno.EINTR:
            raise

    except OSError, e:
      if e[0] != errno.EIO:             # EIO when the shell exits
        raise

  finally:
    reset_mode()
    print >>sys.stderr, '*** Closed the session for "%s".' % ' '.join(argv)


#
# Main
#

def usage():
  print >>sys.stderr, "Usage: %s [-x] [-f] [-c port | program-to-exec]" % progname

xterm = os.environ.has_key('DUALPYXTERM') or 'xterm' in os.environ.get('TERM', '')

try:
  dstport = 0
  force = False
  (opts, argv) = getopt.getopt(
      sys.argv[1:], 'xfc:hv', ['xterm', 'force', 'version', 'help'])
  for (o, a) in opts:
    if o == '-c':
      dstport = a
    elif o in ('-x', '--xterm'):
      xterm = 1
    elif o in ('-f', '--force'):
      force = True
    elif o in ('-v', '--version'):
      print >>sys.stderr, '%s: %s' % (progname, version)
      sys.exit(os.EX_OK)
    elif o in ('-h', '--help'):
      usage()
      sys.exit(os.EX_OK)
except getopt.GetoptError:
  usage()
  sys.exit(os.EX_USAGE)

if not force and os.environ.has_key(ENV_FLAGNAME):
  print >>sys.stderr, "%s: Already in the %s environment" % (progname, progname)
  sys.exit(os.EX_USAGE)

try:
  curses.setupterm()

  if dstport:
    if argv:
      usage()
      sys.exit(os.EX_USAGE)
    try:
      client(int(dstport))
    except ValueError:
      usage()
      sys.exit(os.EX_USAGE)

  else:
    if not argv:
      s = os.environ.get('SUDO_USER') or \
          os.environ.get('LOGNAME') or \
          os.environ.get('USER')
      try:
        s = s and pwd.getpwnam(s)[6]
      except KeyError:
        pass
      try:
        s = s or pwd.getpwuid(os.getuid())[6]
      except KeyError:
        pass
      argv = (s or '/bin/bash',)
    server(argv)

except SystemExit:
  raise

except KeyboardInterrupt:
  print >>sys.stderr, '^C'
  sys.exit(os.EX_OK)

except OSError, e:
  print >>sys.stderr, '%s: %s' % (progname, e[1])
  sys.exit(os.EX_OSERR)

except Exception:
  print >>sys.stderr, '%s: Unknown error' % progname
  sys.exit(os.EX_SOFTWARE)