/reassembly.py
Python | 125 lines | 113 code | 1 blank | 11 comment | 0 complexity | 743e104f310fbf8c037511d8dba87fa1 MD5 | raw file
- from dis import disassemble
- from opcode import HAVE_ARGUMENT, opname, opmap, hasjabs
- import inspect
- from StringIO import StringIO
- import new
- '''
- Doesn't work.
-
- The problem is, I can't capture the local value stack and, more importantly,
- block stack. Which means that `for i in range(10)` can't be properly restored,
- given that range iterator is stored in one of the them.
- '''
- class Saved_frame(object):
- def __init__(self, ip, locals, retval):
- self.ip = ip
- self.locals = locals
- self.retval = retval
-
- def save_frame(retval):
- frame = inspect.currentframe(1)
- ldict = frame.f_locals
- print 'Locals: {!r}'.format(ldict)
- locals = []
- for name in frame.f_code.co_varnames:
- print name, ldict[name]
- locals.append(ldict[name])
- return Saved_frame(frame.f_lasti, locals, retval)
-
- class Instruction(object):
- def __init__(self, addr, op, arg = None):
- self.addr = addr
- self.op = op
- self.arg = arg
- def __repr__(self):
- s = '{:>3}: {}'.format(self.addr, opname[self.op])
- if self.arg is not None:
- s += ' ' + str(self.arg)
- return s
-
- def disassemble_to_list(code):
- lst = []
- n = len(code)
- i = 0
- while i < n:
- c = code[i]
- op = ord(c)
- addr = i
- i += 1
- oparg = None
- if op >= HAVE_ARGUMENT:
- oparg = ord(code[i]) + ord(code[i + 1])*256
- i += 2
- lst.append(Instruction(addr, op, oparg))
- return lst
-
- def assemble_from_list(lst, offset = 0):
- out = StringIO()
- for it in lst:
- out.write(chr(it.op))
- if it.op >= HAVE_ARGUMENT:
- arg = it.arg
- if it.op in hasjabs:
- arg += offset
- out.write(chr(arg & 0xFF))
- out.write(chr(arg >> 8))
- return out.getvalue()
-
-
- def patch(f):
- # new.function(code, globals[, name[, argdefs[, closure]]])
- # new.code(argcount, nlocals, stacksize, flags, codestring, constants, names, varnames, filename, name, firstlineno, lnotab)
- code = f.__code__
- lst = disassemble_to_list(code.co_code)
- local_cnt = code.co_nlocals
-
- for it in lst:
- if it.op == opmap['RETURN_VALUE']:
- return_addr = it.addr + 1
- break
-
- patch = [
- Instruction(0, opmap['LOAD_FAST'], 0),
- Instruction(0, opmap['UNPACK_SEQUENCE'], local_cnt),]
-
- for i in range(local_cnt):
- patch.append(Instruction(0, opmap['STORE_FAST'], i))
-
- assembled = assemble_from_list(patch)
- patch_length = len(assembled) + 3
-
- patch = [Instruction(0, opmap['JUMP_ABSOLUTE'], return_addr + patch_length)]
- assembled += assemble_from_list(patch)
-
- assert len(assembled) == patch_length
- assembled += assemble_from_list(lst, patch_length)
-
- patched_code = new.code(1, code.co_nlocals, code.co_stacksize,
- code.co_flags, assembled, code.co_consts, code.co_names,
- code.co_varnames, code.co_filename, code.co_name,
- code.co_firstlineno, code.co_lnotab)
- patched_function = new.function(patched_code, f.func_globals, 'generator', None, f.func_closure)
- return patched_function
-
- def facs(step = 1):
- acc = 1
- n = 0
- while True:
- # for _ in range(10):
- n += step
- acc *= n
- return save_frame(acc)
- acc += 0 # to prevent the compiler from removing the jump back
-
- facs2 = patch(facs)
-
- print disassemble(facs.__code__)
- print disassemble(facs2.__code__)
- # print disassemble(facs_patched.__code__)
- print 'start'
- frame = facs()
- print frame.__dict__
- frame = facs2(frame.locals)
- print frame.__dict__
- frame = facs2(frame.locals)
- print frame.__dict__