Function smash!

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 = "<smash>"
+    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)