root/projects/AsynCluster/trunk/console

#!/usr/bin/env python
#
# AsynCluster: Master
# A cluster management server based on Twisted's Perspective Broker. Dispatches
# cluster jobs and regulates when and how much each user can use his account on
# any of the cluster node workstations.
#
# Copyright (C) 2006-2007 by Edwin A. Suominen, http://www.eepatents.com
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
# 
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the file COPYING for more details.
# 
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 51
# Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA

"""
An interface console to the controller.
"""

import os, sys, re
from textwrap import wrap
from optparse import OptionParser
from signal import signal, SIGWINCH
from fcntl import ioctl
from tty import TIOCGWINSZ
from struct import unpack

import configobj

from twisted.python.failure import Failure
from twisted.internet import reactor, defer
from twisted.spread import pb

from twisted.conch import stdio
from twisted.conch.insults.insults import TerminalProtocol, privateModes
from twisted.conch.insults.window import TopWindow, VBox, TextInput, TextOutput


CONNECT_DELAY = 2.0
CONFIG_PATH = "/etc/asyncluster.conf"


class Client(object):
    """
    I connect to the master TCP server via the UNIX socket control interface.
    """
    root = None
    
    def __init__(self, socket):
        if not os.path.exists(socket):
            raise RuntimeError(
                "No UNIX socket available at '%s'" % socket)
        self.socket = socket
    
    def connect(self):
        """
        Makes the UNIX socket connection, storing a remote reference to the
        server's control root object if the connection is successful. Returns a
        deferred that fires C{True} if so, and I am thus ready to accept
        commands as a result, or C{False} otherwise.
        """
        def gotAnswer(answer):
            if pb.IUnjellyable.providedBy(answer):
                self.root = answer
                return True
            self.connector.disconnect()
            return False

        factory = pb.PBClientFactory()
        self.connector = reactor.connectUNIX(self.socket, factory)
        return factory.getRootObject().addBoth(gotAnswer)

    def disconnect(self):
        """
        Disconnects from the master TCP server, returning a deferred that fires
        when the disconnection is complete. Before the TCP disconnection
        occurs, any jobs that are running are allowed to finish and any active
        session is ended.
        """
        self.root = None
        self.connector.disconnect()

    def _jobCode(self, filePath):
        """
        """
        fh = open(filePath)
        result = fh.read()
        fh.close()
        return result

    def oops(self, msg):
        return defer.succeed("ERROR: %s" % msg)

    def command(self, cmd, *args):
        """
        """
        caller = getattr(self.root, 'callRemote', None)
        if not callable(caller):
            return self.oops("No control connection available")

        if cmd == 'user':
            cmd = 'userAction'
        elif cmd == 'wall':
            args = (" ".join(args),)
        elif cmd == 'resetup':
            srcDir = args[0]
            if not os.path.isdir(srcDir):
                return self.oops("Source directory '%s' not found" % srcDir)
        return caller(cmd, *args)


class History(object):
    """
    Self-contained input history representation.

    @type beforeLines: C{list} of C{str}
    @ivar beforeLines: The lines in this history which come before the current position.

    @type afterLines: C{list} of C{str}
    @ivar afterLines: The lines in this history which come after the current position.
    """
    def __init__(self, lines=None):
        if lines is None:
            lines = []
        self.beforeLines = lines
        self.afterLines = []

    def nextLine(self):
        """
        Advance the position by one and return the line there, or an empty
        string if there is no next line.
        """
        if self.afterLines:
            self.beforeLines.append(self.afterLines.pop(0))
            if self.afterLines:
                return self.afterLines[0]
            return ""
        return ""

    def previousLine(self):
        """
        Rewind the position by one and return the line there, or an empty
        string if there is no previous line.
        """
        if self.beforeLines:
            self.afterLines.insert(0, self.beforeLines.pop())
            return self.afterLines[0]
        return ""

    def allLines(self):
        """
        Return a list of all lines in this history object.
        """
        return self.beforeLines + self.afterLines

    def addLine(self, line):
        """
        Add a new line to the end of this history object.
        """
        if self.afterLines:
            self.afterLines.append(line)
        else:
            self.beforeLines.append(line)

    def resetPosition(self):
        """
        Set the position in the input history to the end.
        """
        self.beforeLines.extend(self.afterLines)
        self.afterLines = []


class LineInputWidget(TextInput):
    """
    Single-line input area with history and function keys.

    @ivar previousKeystroke: A reference to the most recently received
    keystroke, updated after each keystroke is processed.

    @ivar killRing: A list of killed strings, in order of oldest to newest.

    @type savedBuffer: C{NoneType} or C{str}
    @ivar savedBuffer: The string in the edit buffer at the time a history
    traversal command was first invoked, or C{None} if the history is not
    currently being traversed.

    """
    previousKeystroke = None
    savedBuffer = None

    def __init__(self, maxWidth, onSubmit):
        self._realSubmit = onSubmit
        self.killRing = []
        self.setInputHistory(History())
        super(LineInputWidget, self).__init__(maxWidth, self._onSubmit)

    def setInputHistory(self, history):
        """
        Set the complete input history to the given history object.
        """
        self.inputHistory = history

    def getInputHistory(self):
        """
        Retrieve a list of lines representing the current input history.
        """
        return self.inputHistory.allLines()

    def _onSubmit(self, line):
        """
        Clear the current buffer and call the submit handler specified when
        this widget was created.
        """
        if line:
            self.inputHistory.addLine(line)
            self.inputHistory.resetPosition()
            self.setText('')
        self._realSubmit(line)

    def func_HOME(self, modifier):
        """
        Handle the home function key by repositioning the cursor at the
        beginning of the input area.
        """
        self.cursor = 0

    def func_CTRL_a(self):
        """
        Handle C-a in the same way as the home function key.
        """
        return self.func_HOME(None)

    def func_CTRL_f(self):
        """
        Handle C-f to move the cursor forward one position.
        """
        self.cursor = min(self.cursor + 1, len(self.buffer))

    def func_CTRL_b(self):
        """
        Handle C-b to move the cursor forward one position.
        """
        self.cursor = max(self.cursor - 1, 0)

    def func_END(self, modifier):
        """
        Handle the end function key by repositioning the cursor just past the
        end of the text in the input area.
        """
        self.cursor = len(self.buffer)

    def func_CTRL_e(self):
        """
        Handle C-e in the same way as the end function key.
        """
        return self.func_END(None)

    def func_ALT_b(self):
        """
        Handle M-b by moving the cursor to the beginning of the word under the
        current cursor position.  Do nothing at the beginning of the line.
        Words are considered non-whitespace characters delimited by whitespace
        characters.
        """
        while self.cursor > 0 and self.buffer[self.cursor - 1].isspace():
            self.cursor -= 1
        while self.cursor > 0 and not self.buffer[self.cursor - 1].isspace():
            self.cursor -= 1

    def func_ALT_f(self):
        """
        Handle M-f by moving the cursor to just after the end of the word under
        the current cursor position.  Do nothing at the end of the line.  Words
        are considered non-whitespace characters delimited by whitespace
        characters.
        """
        n = len(self.buffer)
        while self.cursor < n and self.buffer[self.cursor].isspace():
            self.cursor += 1
        while self.cursor < n and not self.buffer[self.cursor].isspace():
            self.cursor += 1

    def func_CTRL_k(self):
        """
        Handle C-k by truncating the line from the character beneath the cursor
        and adding the removed text to the kill ring.
        """
        chopped = self.buffer[self.cursor:]
        self.buffer = self.buffer[:self.cursor]
        if chopped:
            self.killRing.append(chopped)

    def func_CTRL_y(self):
        """
        Handle C-y by inserting an element from the kill ring at the current
        cursor position, moving the cursor to the end of the inserted text.
        """
        if self.killRing:
            insert = self.killRing[-1]
            self.buffer = self.buffer[:self.cursor] + insert + self.buffer[self.cursor:]
            self.cursor += len(insert)

    def func_ALT_y(self):
        """
        Handle M-y by cycling the kill ring and replacing the previously yanked
        text with the new final element in the ring.
        """
        if self.previousKeystroke in (('\x19', None), ('y', ServerProtocol.ALT)): # C-y and M-y
            previous = self.killRing.pop()
            self.killRing.insert(0, previous)
            next = self.killRing[-1]

            self.cursor -= len(previous)
            self.buffer  = self.buffer[:self.cursor] \
                           + next + self.buffer[self.cursor + len(previous):]
            self.cursor += len(next)

    def func_CTRL_p(self):
        """
        Handle C-p to swap the current input buffer with the previous line from
        input history.
        """
        if not self.inputHistory.afterLines:
            # Going from normal editing to history traversal - save the edit
            # buffer.
            self.savedBuffer = self.buffer
        previousLine = self.inputHistory.previousLine()
        if previousLine:
            self.buffer = previousLine

    def func_CTRL_n(self):
        """
        Handle C-n to swap the current input buffer with the next line from
        input history.
        """
        nextLine = self.inputHistory.nextLine()
        if nextLine:
            self.buffer = nextLine
        else:
            if self.savedBuffer is not None:
                self.buffer = self.savedBuffer
                self.savedBuffer = None

    def keystrokeReceived(self, keyID, modifier):
        """
        Override the inherited behavior to track whether either the cursor
        position or buffer contents change and automatically request a repaint
        if either does.
        """
        buffer = self.buffer
        cursor = self.cursor
        super(LineInputWidget, self).keystrokeReceived(keyID, modifier)
        self.previousKeystroke = (keyID, modifier)
        if self.buffer != buffer or self.cursor != cursor:
            self.repaint()

    def characterReceived(self, keyID, modifier):
        """
        Handle a single non-function key, possibly with a modifier.

        If there is no modifier, let the super class handle this.  Otherwise,
        dispatch to a function for the specific key and modifier present.
        """
        if modifier is not None:
            f = getattr(self, 'func_' + modifier.name + '_' + keyID, None)
            if f is not None:
                f()
        elif ord(keyID) <= 26 and keyID != '\r':
            f = getattr(self, 'func_CTRL_' + chr(ord(keyID) + ord('a') - 1), None)
            if f is not None:
                f()
        else:
            super(LineInputWidget, self).characterReceived(keyID, modifier)


class OutputWidget(TextOutput):
    def __init__(self, size=None):
        super(OutputWidget, self).__init__(size)
        self.messages = []

    def formatMessage(self, s, width):
        return wrap(s, width=width, subsequent_indent="  ")

    def addMessage(self, message):
        self.messages.append(message)
        self.repaint()

    def render(self, width, height, terminal):
        output = []
        for i in xrange(len(self.messages) - 1, -1, -1):
            output[:0] = self.formatMessage(self.messages[i], width - 2)
            if len(output) >= height:
                break
        if len(output) < height:
            output[:0] = [''] * (height - len(output))
        for n, L in enumerate(output):
            terminal.cursorPosition(0, n)
            terminal.write(L + ' ' * (width - len(L)))


class UserInterface(TerminalProtocol):
    """
    Set up an input area and an output area for the console.
    """
    width, height = 80, 24
    
    reactor = reactor
    ready = False

    def _get_ID(self):
        self._ID = (getattr(self, '_ID', 0) + 1) % 99
        return self._ID
    ID = property(_get_ID)

    def _get_config(self):
        return configobj.ConfigObj(CONFIG_PATH)
    config = property(_get_config)

    def connectionMade(self):
        """
        Called when the terminal is connected.
        """
        super(UserInterface, self).connectionMade()
        self.terminal.eraseDisplay()
        self.terminal.resetPrivateModes([privateModes.CURSOR_MODE])

        self.rootWidget = TopWindow(
            self._painter, lambda f: self.reactor.callLater(0, f))
        self.rootWidget.reactor = self.reactor
        vbox = VBox()
        vbox.addChild(OutputWidget())
        vbox.addChild(LineInputWidget(self.width-2, self.parseInputLine))
        self.rootWidget.addChild(vbox)
        self.reactor.callLater(CONNECT_DELAY, self.startup)

    def _painter(self):
        self.rootWidget.draw(self.width, self.height, self.terminal)

    def addOutputMessage(self, msg):
        return self.rootWidget.children[0].children[0].addMessage(msg)

    def startup(self):
        def doneTrying(success):
            self.ready = success
            if success:
                msg = "+++ Connected to Master Control Server +++"
            else:
                msg = "--- Couldn't connect to Master Control Server! ---"
            self.addOutputMessage(msg)

        self.client = Client(self.config['common']['socket'])
        return self.client.connect().addCallback(doneTrying)

    def shutdown(self):
        def disconnected(null):
            self.terminal.setPrivateModes([privateModes.CURSOR_MODE])
            self.terminal.loseConnection()

        return self.client.disconnect().addCallback(disconnected)

    def parseInputLine(self, line):
        """
        """
        def gotResult(result, ID):
            self.addOutputMessage("<- %02d: %s" % (ID, str(result)))

        def gotFailure(failure, ID):
            msg = failure.getTraceback()
            self.addOutputMessage("<- %02d: ERROR: %s" % (ID, msg))

        tokens = line.split()
        if tokens[0] in ('q', 'quit'):
            return self.shutdown()

        ID = self.ID
        self.addOutputMessage("%02d -> %s" % (ID, line))
        
        d = self.client.command(*tokens)
        d.addCallback(gotResult, ID)
        d.addErrback(gotFailure, ID)

    def keystrokeReceived(self, keyID, modifier):
        self.rootWidget.keystrokeReceived(keyID, modifier)

    def terminalSize(self, width, height):
        self.width = width
        self.height = height
        self._painter()


class CommandLineUserInterface(UserInterface):
    """
    """
    def connectionMade(self):
        signal(SIGWINCH, self.windowChanged)
        winSize = self.getWindowSize()
        self.width = winSize[0]
        self.height = winSize[1]
        super(CommandLineUserInterface, self).connectionMade()

    def connectionLost(self, reason):
        reactor.stop()

    # XXX Should be part of runWithProtocol
    def getWindowSize(self):
        winsz = ioctl(0, TIOCGWINSZ, '12345678')
        winSize = unpack('4H', winsz)
        newSize = winSize[1], winSize[0], winSize[3], winSize[2]
        return newSize

    def windowChanged(self, signum, frame):
        winSize = self.getWindowSize()
        self.terminalSize(winSize[0], winSize[1])


if __name__ == '__main__':
    stdio.runWithProtocol(CommandLineUserInterface)