commit 545a3ca1947c0048db12d90d1a34b897882d5b2f Author: Ben Kurtovic Date: Mon Apr 16 21:41:47 2012 -0400 Function smash! diff --git a/func_smash.py b/func_smash.py new file mode 100644 index 0000000..61323e7 --- /dev/null +++ b/func_smash.py @@ -0,0 +1,119 @@ +import dis +import opcode +import random +import types + +############## CODE ########################################################### + +MARKOV_START = -1 +MARKOV_END = -2 + +def make_chain(funcs): + chain = {} + for func in funcs: + co = func.func_code + code = co.co_code + n = len(code) + i = 0 + lastop = MARKOV_START + while i < n: + c = code[i] + op = ord(c) + i += 1 + if op >= opcode.HAVE_ARGUMENT: + oparg = ord(code[i]) + ord(code[i + 1]) * 256 + i += 2 + if op in opcode.hasconst: + arg = co.co_consts[oparg] + elif op in opcode.haslocal: + arg = co.co_varnames[oparg] + else: + raise NotImplementedError(op, opcode.opname[op]) + else: + arg = None + _chain_append(chain, lastop, (op, arg)) + lastop = op + _chain_append(chain, op, MARKOV_END) + return chain + +def make_function(chain, name, argcount=1): + codes, constants, varnames = _make_codes(chain) + nlocals = len(varnames) + argcount + stacksize = 1024 # High limit? + flags = 0 # Denotes funcs with *args and/or **kwargs; nothing for now + codestring = "".join([chr(code) for code in codes]) + names = () + filename = "" + firstlineno = 1 + lnotab = "" + + code = types.CodeType(argcount, nlocals, stacksize, flags, codestring, + constants, names, varnames, filename, name, + firstlineno, lnotab) + func = types.FunctionType(code, globals(), name) + return func + +def _chain_append(chain, first, second): + try: + chain[first].append(second) + except KeyError: + chain[first] = [second] + +def _make_codes(chain): + codes = [] + code = random.choice(chain[MARKOV_START]) + constants, varnames = [], [] + while 1: + if code == MARKOV_END: + break + op, arg = code + codes.append(op) + if op >= opcode.HAVE_ARGUMENT: + if op in opcode.hasconst: + args = constants + elif op in opcode.haslocal: + args = varnames + else: + raise NotImplementedError(op, opcode.opname[op]) + if arg not in args: + args.append(arg) + codes.append(args.index(arg)) + codes.append(0) + code = random.choice(chain[op]) + return codes, tuple(constants), tuple(varnames) + +############## FUNCTION CORPUS ################################################ + +def f1(a): + b = a + 7.0 + return b + +def f2(a): + b = a - 5.0 + return b + +def f3(a): + b = a * 3.0 + return b + +def f4(a): + b = a / 2.0 + return b + +corpus = [f1, f2, f3, f4] + +############## DEMO ########################################################### + +if __name__ == "__main__": + print "USING FUNCTION CORPUS:", [func.__name__ for func in corpus] + chain = make_chain(corpus) + print "OPCODE MARKOV CHAIN:", chain + print + + func = make_function(chain, "func") + print "SMASHED FUNCTION CODE:", [ord(c) for c in func.func_code.co_code] + print "FUNCTION DISASSEMBLY:" + dis.dis(func) + print + n = 12.0 + print "func(%s) =" % n, func(n)