Source

prise / src / state.py

#
# prise - state.py
#
# Copyright (c) 2012, Dan Boitnott All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
# 
# Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer. Redistributions in
# binary form must reproduce the above copyright notice, this list of
# conditions and the following disclaimer in the documentation and/or
# other materials provided with the distribution. THIS SOFTWARE IS
# PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# This module maintains global state for all other modules and provides the
# core methods which control shell interaction.

import sys
import os
import socket
import threading
import glob

import util
import output
import hotkey
import scripts

shl = None
currentLine = ''
runDir = os.path.join(os.getcwd(), ".prise", "runs", socket.gethostname() + "-" + util.timestamp())

# Scripting
script = None
scriptIndex = 0
scriptLock = threading.RLock()
scriptPath = None

# Create the run directory
os.makedirs(runDir)

# Open the run log
runLog = open(os.path.join(runDir, "run.log"), "a")

# Open output multiplexer
multiOutput = output.MultiOutput()
multiOutput.addFile(runLog)

# Command Logging
curCmdLog = None
curCmdLogLock = threading.RLock()

# Web interface
webServer = None

def inputFilter(str):
    if str not in hotkeys:
        return str

    hotkeys[str]()

    return ''

def outputFilter(str):
    global currentLine, waitingOutput, scriptIndex
    # Attempt to detect lines to maintain current line buffer
    if str == '' or str[-1] in ('\n', '\r'):
        currentLine = ''
    else:
        currentLine = str
        if atPrompt():
            closeCmdLog()

            # Output any comments and move to an actual command
            scriptLock.acquire()
            try:
                while scriptIndex < len(script) and type(script[scriptIndex]) is scripts.CommentTask:
                    script[scriptIndex].do()
                    scriptIndex += 1
            finally:
            	scriptLock.release()

    return str

def atPrompt():
    return currentLine.startswith("prise:")

def closeCmdLog():
    global curCmdLog
    curCmdLogLock.acquire()
    try:
        if curCmdLog:
            multiOutput.removeFile(curCmdLog)
            print >>curCmdLog, "\n", util.RULE
            print >>curCmdLog, "Finished:", util.prettyTimestamp()
            curCmdLog.close()
            curCmdLog = None
    finally:
        curCmdLogLock.release()

def getCommandOutputPathForKey(key):
    ret = os.path.join(runDir, "cmd.%s.log" % key)
    if not util.isSubfile(runDir, ret):
        return None
    return ret

def getCommandOutputPath(cmd, ts):
    return getCommandOutputPathForKey(util.cmdHash(cmd.strip()) + "." + ts)

def logCommand(cmd):
    curCmdLogLock.acquire()
    try:
        global curCmdLog
        if curCmdLog:
            closeCmdLog()

        cmd = cmd.strip()
        curCmdLog = open(getCommandOutputPath(cmd, util.timestamp()), "w")
        util.writeFileHeader(curCmdLog, { "Command" : cmd,
                                          "Started" : util.prettyTimestamp(),
                                          "Host" : socket.gethostname()})
        multiOutput.addFile(curCmdLog)
    finally:
        curCmdLogLock.release()

def runCommand(cmd):
    logCommand(cmd)
    sendToShell(cmd + "\n")

def getCommandOutputTimes(cmd, desc = True):
    """Returns the timestamps for each run of a given command."""

    if type(cmd) is scripts.CommandTask:
        # Extract the command from the task
        cmd = cmd.content.strip()
    elif type(cmd) is scripts.CommentTask:
        # Deal gracefully with being asked for output times on comments
        return []
    else:
        cmd = cmd.strip()

    glb = getCommandOutputPath(cmd, '*')
    return sorted([os.path.basename(p).split('.', 2)[2][:-4] for p in glob.glob(glb)], reverse=desc)

def getCommandOutput(cmd, ts):
    """Returns the output of a command run at the given timestamp. ts should exactly match the one returned from
       getCommandOutputTimes()"""

    return output.CommandOutput(getCommandOutputPath(cmd, ts))

def getCommandOutputForKey(key):
    return output.CommandOutput(getCommandOutputPathForKey(key))

def printToUser(msg, reprompt = True):
    # Print to STDERR so that it won't be sent to the shell
    print >>sys.stderr, "\r\n", msg.replace("\n", "\n\r"), "\r"

    if reprompt and atPrompt():
        # Re-print the prompt for the user
        print >>sys.stderr, currentLine,

    sys.stderr.flush()

def sendToShell(str):
    # Send a command to the shell
    shl.send(str)

def shutdown():
    multiOutput.close()
    if webServer:
        webServer.shutdown()

def getLastCmdIdx():
    scriptLock.acquire()
    try:
        ret = scriptIndex - 1
        while ret >= 0 and type(script[ret]) is scripts.CommentTask:
            ret -= 1
        return ret
    finally:
        scriptLock.release()

def runLastCmd():
    scriptLock.acquire()
    try:
        idx = getLastCmdIdx()
        if idx < 0:
            printToUser("You are at the first command.")
            return
        script[idx].do()
    finally:
        scriptLock.release()
    	
def showLastCmd():
    scriptLock.acquire()
    try:
        idx = getLastCmdIdx()
        if idx < 0:
            printToUser("You are at the first command.")
            return
        printToUser("Last: " + script[idx].content)
    finally:
        scriptLock.release()
    	
def checkScriptIndex():
    scriptLock.acquire()
    try:
        global scriptIndex
        if scriptIndex >= len(script):
            printToUser("::: End of script :::")
            scriptIndex = len(script)
            return True
        return False
    finally:
        scriptLock.release()
    	
def incrScriptIndex():
    scriptLock.acquire()
    try:
        global scriptIndex
        scriptIndex += 1
        checkScriptIndex()
    finally:
        scriptLock.release()
    	
def skipNextCmd():
    incrScriptIndex()
    showNextCmd()

def runNextCmd():
    scriptLock.acquire()
    try:
        if checkScriptIndex():
            return
        script[scriptIndex].do()
        incrScriptIndex()
    finally:
        scriptLock.release()

def showNextCmd():
    scriptLock.acquire()
    try:
        if checkScriptIndex():
            return
        printToUser("Next: " + script[scriptIndex].content)
    finally:
        scriptLock.release()

def hotkeyHelp():
    printToUser("""Hotkey Help:
    F1: Print help information
    F2: Print script
    F5: Re-run last command
    F6: Show last command
    F8: Skip next command
    F9: Run next command
    F10: Show next command""")

def showScript():
    out = ''
    for i in range(len(script)):
        if i == scriptIndex:
            tick = '>'
        else:
            tick = ' '
        first = True
        for l in script[i].getLines():
            idxCol = "%3d" % i if first else "   "
            out += "%s %s %s\n" % (idxCol, tick, l)
            first = False
    printToUser(out)

def writeScript(path):
    scriptLock.acquire()
    try:
        fp = open(path, 'w')
        try:
            for task in script:
                for line in task.getLines():
                    print >>fp, line
        finally:
            fp.close()
    finally:
        scriptLock.release()
    	
def backupScript():
    scriptLock.acquire()
    try:
        writeScript(os.path.join(runDir, "script.%s.ise" % util.timestamp()))
    finally:
        scriptLock.release()

def saveScript():
    scriptLock.acquire()
    try:
        backupScript()
        writeScript(scriptPath)
    finally:
        scriptLock.release()


hotkeys = {
    hotkey.KEY_F1 : hotkeyHelp,
    hotkey.KEY_F2 : showScript,
    hotkey.KEY_F5 : runLastCmd,
    hotkey.KEY_F6 : showLastCmd,
    hotkey.KEY_F8 : skipNextCmd,
    hotkey.KEY_F9 : runNextCmd,
    hotkey.KEY_F10 : showNextCmd
}