iesh

#!/usr/bin/env python
# iesh / ie_shell.py - Simple shell for Infinity Engine-based game files
# Copyright (C) 2004-2010 by Jaroslav Benkovsky, <edheldil@users.sf.net>
#
# 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
# GNU General Public License 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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.


import atexit
import os
import os.path
import pdb
import rlcompleter
import readline
import subprocess
import sys
import traceback

import infinity
from infinity import core
from infinity.builtins import *
from infinity.stream import *
from infinity.formats import *

###################################################
IESH_PROFILE_FILE = "~/.iesh_profile"
"""Load and execute this file at startup in interactive mode."""

IESH_HISTORY_FILE = "~/.iesh_history"
"""Command line history"""

IESH_RESTORE_FILE = "~/.iesh_restore"
"""Save keys, strrefs and other core vars to this file for quick restore."""

IESH_NORMAL_PROMPT = 'Cmd: '
IESH_COMPOUND_PROMPT = '.... '
IESH_INDENT = '    '

###################################################
## list of commands to be executed before eventual user's input
commands = []

## current indentation for compound commands
indent = ""

## current compound command if non-empty
compound_command = ""

pager = None
"""Pager process (as a Popen object)"""

quit = False

###################################################
def help_on_shell ():
    print """
    This is basically a python shell, meaning that the commands you enter
    are mostly python commands and statements. Compound commands have
    to be terminated by an empty line.

    Use TAB to complete commands, object, attribute or method names, etc.
    Use UP or DOWN cursor keys to navigate in command history
    Use LEFT or RIGHT cursor keys to edit current command
    Use ^R to search in command history

    Some notable commands:
      load_game ("/home/ed/pst", "CHITIN.KEY", "dialog.tlk")
          Loads key and dialog files from the specified directory.
          The directory parameter is mandatory, the others are optional.
          Most of the other commands assume that these two files are
          already loaded. The loaded objects are stored in core.keys and
          core.strrefs.

      find_objects('demogor2')
      find_objects('demogor2', 'cre')
      find_objects('demogor2', 1009)
          Tries to find all objects that match the name and, optionally, file type
          inside the game's data and the override folder.

      load_object ("data/AGOODY.itm")
      load_object ("AGOODY")
          tries to find and load object from a file or resources.
          Returns format object, so the object type has to be supported

      load_object ("AGOODY").printme ()

      export_object ("GUICG", "GUICG.chu")
          Try to find resource GUICG and export it to file GUICG.chu.
          If the name is not unique, add third parameter,
          e.g. type=0x03ea to restrict the resource type

      find_str ("^(?i)gemrb")
          List strrefs for all strings starting with word GemRB, regardless
          of case


      c = chui.CHUI_Format ("GUISAVE.chu")
      c.read ()
      c.printme ()
           loads file GUISAVE.chu from the current dir, decodes it
           and prints its contents

           key.KEY_Format
           tlk.TLK_Format
           biff.BIFF_Format

           bam.BAM_Format
           chui.CHUI_Format
           cre.CRE_Format   (only parts)
           ids.IDS_Format
           itm.ITM_Format
           spl.SPL_Format   (NOT WORKING!)
           wmap.WMAP_Format

      core.keys.print_bif_record (core.keys.bif_list[0])
      core.strrefs.print_strref_record (core.strrefs.strref_list[0])
           print first records from chitin.key or dialog.tlk file

      help (__name__)
      help (infinity)
      help (infinity.builtins)
           help on various topics

      !shell_cmd
          Runs command shell_cmd in your native shell. Useful to list
          files on disk, for example.

      @filename
          Execute python commands from the named file one by one

     ?name
          Print help on name

      quit
          Exit the shell.

    """


###################################################
def read_command_file (filename):
    """Read commands to be executed from a file."""

    try: fh = open (filename)
    except: return False

    commands.extend (map (lambda s: s.rstrip ("\n\r"), fh.readlines ()))

    return True

###################################################
def read_command (prompt):
    """Return next command to be executed.
    The command is read either from a file or from the command line."""

    global quit

    if quit:
        return 'quit'

    s = ''
    if len (commands) != 0:
        s = commands.pop (0)
    else:
        try:
            if sys.stdin.isatty():
                s = raw_input (prompt)
            else:
                s = sys.stdin.readline()
                if s.endswith("\n"):
                    s = s[:-1]
                elif s == '':
                    quit = True
        except KeyboardInterrupt:
            pass
        except EOFError:
            quit = True

    return s

###################################################
def _readline_init ():
    """Enable history and tab completion for the command line."""

    readline.parse_and_bind ("tab: complete")
    histfile = os.path.expanduser (IESH_HISTORY_FILE)
    try:
        readline.read_history_file (histfile)
    except IOError:
        pass

    atexit.register (readline.write_history_file, histfile)
    #del histfile
    readline.set_pre_input_hook (_readline_hook)

###################################################
def _readline_hook ():
    readline.insert_text (indent)
    readline.redisplay ()

###################################################
def save (name = ''):
    """Saves some core variables (especially keys and strrefs), so that they
    can be later loaded faster than from IE data files. if `name' is specified,
    it's added to filename where the data is stored, thus allowing more
    save files than just one."""
    if name != '':
        name = '-' + name
    save_state (os.path.expanduser (IESH_RESTORE_FILE + name))

###################################################
def restore (name = ''):
    """Loads some core variables (especially keys and strrefs) saved
    by save() function, much faster than from IE data files. if `name'
    is specified, it's added to filename where the data is stored, thus
    allowing more save files than just one."""
    if name != '':
        name = '-' + name
    restore_state (os.path.expanduser (IESH_RESTORE_FILE + name))

###################################################
###################################################

#
# If the program is run with a file parameter, the contents of
# the file are executed and the program quits.
# If it's run without arguments, the program loads and executes
# profile file and starts interactive command prompt
#

if len (sys.argv) > 1:
    sys.stdin = open(sys.argv[1])
elif sys.stdin.isatty():
    _readline_init ()
    read_command_file (os.path.expanduser (IESH_PROFILE_FILE))
    print "\n\nType `help' to print a short help, `quit' or ^D to exit the shell\n"

#
# Main command processing loop
#

while True:
    if compound_command == '':
        prompt = IESH_NORMAL_PROMPT
    else:
        prompt = IESH_COMPOUND_PROMPT

    current_command = read_command (prompt)

    indent = current_command[0:(len (current_command) - len (current_command.lstrip ()))]

    if current_command.endswith (":"):
        indent = indent + IESH_INDENT

    if (current_command != '' and compound_command != '') or current_command.endswith (":"):
        compound_command = compound_command + "\n" + current_command
        continue

    # zero-length input ends compound commands
    if current_command == '' and compound_command != '':
        current_command = compound_command
        compound_command = ''


    if current_command in ('q', 'quit'):
        break

    if current_command in ('?', 'help'):
        help_on_shell ()
        continue

    if current_command.startswith ('?'):
        eval ('help(%s)' %current_command[1:])
        continue

    if current_command.endswith ('?'):
        eval ('help(%s)' %current_command[:-1])
        continue

    if current_command.startswith ('!'):
        os.system (current_command[1:])
        continue

    if current_command.startswith ('@'):
        read_command_file (current_command[1:])
        continue

    # If `pager' option is set, run the `pager' and pipe command output to it
    #   That makes problems with e.g. help(), which runs pager as well, so
    #   try to avoid the situation. Ugly hack, there might be other such commands
    #    and moreover it's too easily fooled.

    save_stdout = sys.stdout
    save_stderr = sys.stdout

    if sys.stdin.isatty() and core.get_option ('pager') and not current_command.startswith ('help') and not current_command.startswith ('load_game') and not current_command.startswith ('~'):
        pager = subprocess.Popen (core.get_option ('pager'),  shell = True,  stdin=subprocess.PIPE)
        sys.stdout = pager.stdin
        sys.stderr = pager.stdin

    try:
        if not current_command.startswith ('~'):
            exec (current_command)
        else:
            pdb.run (current_command[1:])

    except IOError:
        try:
            traceback.print_exc()
            print >> sys.stderr
        except IOError:
            # second IOError probably means that the pager exited,
            #   so even the first exception is not interesting
            if save_stderr:
                print >> save_stderr

    except KeyboardInterrupt:
        pass

    except:
        traceback.print_exc ()
        print >> sys.stderr

    if pager is not None:
        pager.stdin.close ()
        pager.wait ()
        sys.stdout = save_stdout
        sys.stderr = save_stderr
        pager = None


###################################################
# End of file iesh