import code
import dis
import modjelly
import new
import sys
from traceback import format_exc
from types import ModuleType
#from c.l.p
POP_TOP = dis.opname.index("POP_TOP")
PRINT_EXPR = dis.opname.index("PRINT_EXPR")
RETURN_VALUE = dis.opname.index("RETURN_VALUE")
def super_eval(expr, _globals=None, _locals=None):
    if _globals is None: _globals = globals()
    if _locals is None: _locals = _globals
    code = compile(expr, "", "single")

    codestring = code.co_code
    newcode = ""

    last_i = 0
    i = 0
    while i < len(codestring):
        op = ord(codestring[i])
        if op > dis.HAVE_ARGUMENT:
            newcode += codestring[i:i+3]
            i += 3
        elif op == PRINT_EXPR:
            last_i = i
            newcode += chr(POP_TOP)
            i += 1
        else:
            newcode += chr(op)
            i += 1

    if last_i > 0:
        newcode = newcode[:last_i] + chr(RETURN_VALUE) + newcode[last_i+1:]

    code = new.code(code.co_argcount, code.co_nlocals, code.co_stacksize,
                    code.co_flags, newcode, code.co_consts, code.co_names,
                    code.co_varnames, code.co_filename, code.co_name,
                    code.co_firstlineno, code.co_lnotab)

    return eval(code,_globals,_locals) 
allowed = ['array','base64','binascii','bisect','calendar','cmath',
                'collections','copy','datetime','decimal','difflib',
                'heapq','itertools','math','operator','random','re','shlex',
                'string','textwrap','this','time']

def whitelist_import(allowed_modules, _globals, _locals, __import=__import__):
    def restricted_import(name, globals=_globals, locals=_locals, fromlist=None):
	if name not in allowed_modules:
	    raise ImportError,"'%s' module is restricted" % name
	else:
            module = __import(name,globals,locals, fromlist or [])
            if hasattr(module,'__builtins__'):#Not foolproof
                delattr(module,'__builtins__')
	return module
    return restricted_import

def disable(name):
    def disabled(*args,**kwargs):
        raise NotImplementedError,'function %s is disabled for security purposes' % name
    return disabled
def step((buff, prompt, namespace), source, ready='--> ', waiting='... ',maxtrace=0):
    if not source and not buff:
            return ([buff, prompt, namespace],'')
    buff += source
    try:
        compiled = code.compile_command(buff)
    except (SyntaxError,OverflowError,ValueError):
        header,error = format_exc(maxtrace).split('\n',1)
        response = '\n    '.join([header]+error.split('\n'))
        return (['', ready, namespace],response)
    if compiled is None:
        return ([buff+'\n', waiting, namespace],'')
    try:
        echo = super_eval(buff, namespace, namespace)
        if echo is not None:
            response = repr(echo)
        else:
            response = ''
    except:
        header,error = format_exc(maxtrace).split('\n',1)
        response = '\n    '.join([header]+error.split('\n'))
    return (['',ready,namespace],response)
disabled = ['file','open','reload','vars','globals','compile','execfile']
def safe_builtins(disabled_funcs=None, __builtins=__builtins__):
    builtins = ModuleType('__builtins__',__builtins.__doc__)
    builtins.__dict__.update(__builtins.__dict__)
    for name in (disabled_funcs or disabled):
        setattr(builtins,name,disable(name))
    return builtins
def safe_namespace(disabled_funcs=None,allowed_modules=None,__builtins=__builtins__):
    builtins = safe_builtins(disabled_funcs,__builtins)
    namespace = {'__builtins__':builtins,'__name__':'__bastille__'}
    builtins.__import__ = whitelist_import(allowed_modules or allowed, namespace, namespace, __import__)
##    builtins['globals'] = lambda:namespace #Probably shouldn't let them touch it. Maybe a dictproxy...
    return namespace
def interact(banner=None,disabled_funcs=None, allowed_modules=None,__builtins=__builtins__,ready='>>> ',waiting='... '):
    dungeon = safe_namespace(disabled_funcs, allowed_modules)
    if banner: print banner
    state = ('',ready,dungeon)
    while True:
        state, response = step(state, raw_input(state[1]),ready=ready,waiting=waiting)
        if response: print response
def interact2(banner=None,disabled_funcs=None, allowed_modules=None,__builtins=__builtins__,ready='>>> ',waiting='... '):
    dungeon = safe_namespace(disabled_funcs, allowed_modules)
    if banner: print banner
    state = ('',ready,dungeon)
    while True:
        state, response = step(state, raw_input(state[1]),ready=ready,waiting=waiting)
        if response: print response
        state[2]['__builtins__'] = None
        try:
            j = modjelly.jelly(state)
        except modjelly.JellyingError:
            print format_exc()
        else:
            try:
                u = modjelly.unjelly(j)
            except modjelly.UnjellyingError:
                print format_exc()
            else:
                state = u
                state[2]['__builtins__'] = safe_builtins()
            
        
if __name__ == '__main__':
    try:
##        interact(banner='Bastille Test Environment',ready='>>> ',waiting='... ')
        interact2(banner='Bastille Test Environment',ready='>>> ',waiting='... ')
    except (KeyboardInterrupt,EOFError):
        pass