Browse Source

Implementing prettify in func_smash

master
Ben Kurtovic 10 years ago
parent
commit
2c1b1ba24c
2 changed files with 20 additions and 7 deletions
  1. +4
    -0
      func_smash.py
  2. +16
    -7
      prettify.py

+ 4
- 0
func_smash.py View File

@@ -7,6 +7,8 @@ import re
import sys
import types

import prettify

OPMAP = opcode.opmap
OP_HASBUILD = [OPMAP[n] for n in ("BUILD_TUPLE", "BUILD_LIST", "BUILD_MAP",
"BUILD_SET")]
@@ -192,6 +194,8 @@ def _demo(corpus, arg=12.0):
print "Using {0}-function corpus.".format(len(corpus))
print "Smashed function disassembly:"
print_function(func)
print "\nFunction as python code:"
prettify.prettify_function(func, indent=4)

if not len(sys.argv) > 2 or "n" not in "".join(sys.argv[2:]):
print


+ 16
- 7
prettify.py View File

@@ -10,10 +10,12 @@ OP_HASBINARY = {"BINARY_ADD": "+", "BINARY_SUBTRACT": "-",
OP_HASBUILD = {"BUILD_TUPLE": ("(", ")"), "BUILD_LIST": ("[", "]"),
"BUILD_SET": ("{", "}")}

TAB = 4

def prettify_function(func, indent=0):
args = _get_func_args(func)
_print(indent, "def {0}({1}):".format(func.func_name, args))
prettify_code(func.__code__, indent=indent+4)
prettify_code(func.__code__, indent=indent+TAB)

def prettify_code(codeobj, indent=0):
codes = []
@@ -30,11 +32,16 @@ def prettify_code(codeobj, indent=0):
else:
codes.append((code, None))

skip_next_line = False
lines = _parse_codestring(codes)
for i, line in enumerate(lines):
added_indent, code = line[0], line[1:]
if code[-1].endswith(":") and lines[i+1][0] <= added_indent:
if skip_next_line:
skip_next_line = False
continue
added_indent, code = line[0], line[1:]
if code[0] == "else:" and lines[i+2][0] >= added_indent:
# Remove the automatic "pass" after each "if" as it's not needed
skip_next_line = True
_print(indent + added_indent, *code)

def _parse_codestring(codes):
@@ -54,12 +61,13 @@ def _parse_codestring(codes):
i += 1
for x in block_dedent_at:
if i >= x:
indent -= 4
indent -= TAB
block_dedent_at.remove(x)
for x in block_else_at:
if i >= x:
lines.append((indent, "else:"))
indent += 4
indent += TAB
lines.append((indent, "pass"))
block_else_at.remove(x)
_print(indent, i, opname, arg, debug=True)
if opname in OP_HASLOAD:
@@ -118,10 +126,11 @@ def _parse_codestring(codes):
_push(stack, compare)
elif opname == "POP_JUMP_IF_FALSE":
test = _pop(stack)
lines.append((indent, "if {0}:".format(test)))
block_dedent_at.append(arg)
block_else_at.append(arg)
indent += 4
lines.append((indent, "if {0}:".format(test)))
indent += TAB
lines.append((indent, "pass".format(test)))
elif opname == "JUMP_ABSOLUTE":
block_dedent_at.append(i)
elif opname == "JUMP_FORWARD":


Loading…
Cancel
Save