def draw_block_cfg(func, img_file=None):
from xdis import PYTHON_VERSION, IS_PYPY
from control_flow.bb import basic_blocks
from control_flow.cfg import ControlFlowGraph
import dis
import os
name = "some_name"
bb_mgr = basic_blocks(PYTHON_VERSION, IS_PYPY, func)
# dis.dis(func)
cfg = ControlFlowGraph(bb_mgr)
dot_path = '/tmp/flow-%s.dot' % name
if not img_file:
png_path = '/tmp/flow-%s.png' % name
else:
png_path = img_file
with open(dot_path, 'w') as f:
f.write(cfg.graph.to_dot(False))
# print("%s written" % dot_path)
os.system("dot -Tpng %s > %s" % (dot_path, png_path))
if not img_file:
subprocess.run(["xdg-open", png_path])
def test_basic():
for fn in (two_basic_blocks, if_else_blocks):
if debug:
print(fn.__name__)
dis.dis(fn)
print()
bb_mgr = basic_blocks(fn)
check_blocks(bb_mgr.bb_list)
def test_basic():
for fn in (two_basic_blocks, if_else_blocks):
if debug:
print(fn.__name__)
dis.dis(fn)
print()
bb_mgr = basic_blocks(fn)
cfg = ControlFlowGraph(bb_mgr)
if debug:
write_dot(fn.__name__, "/tmp/test_cfg-", cfg.graph, write_png=True)
check_cfg(fn, cfg)
def doit(fn, name=None):
if not name:
name = fn.__name__
print(name)
bb_mgr = basic_blocks(PYTHON_VERSION, IS_PYPY, fn)
for bb in bb_mgr.bb_list:
print("\t", bb)
dis.dis(fn)
cfg = ControlFlowGraph(bb_mgr)
dot_path = '/tmp/flow-%s.dot' % name
png_path = '/tmp/flow-%s.png' % name
open(dot_path, 'w').write(cfg.graph.to_dot(False))
print("%s written" % dot_path)
os.system("dot -Tpng %s > %s" % (dot_path, png_path))
try:
dt = DominatorTree(cfg)
cfg.dom_tree = dt.tree(False)
dfs_forest(cfg.dom_tree, False)
build_dom_set(cfg.dom_tree, False)
dot_path = '/tmp/flow-dom-%s.dot' % name
png_path = '/tmp/flow-dom-%s.png' % name
open(dot_path, 'w').write(cfg.dom_tree.to_dot())
print("%s written" % dot_path)
os.system("dot -Tpng %s > %s" % (dot_path, png_path))
print('*' * 30)
cfg.pdom_tree = dt.tree(True)
dfs_forest(cfg.pdom_tree, True)
build_dom_set(cfg.pdom_tree, True)
dot_path = '/tmp/flow-pdom-%s.dot' % name
png_path = '/tmp/flow-pdom-%s.png' % name
open(dot_path, 'w').write(cfg.pdom_tree.to_dot())
print("%s written" % dot_path)
os.system("dot -Tpng %s > %s" % (dot_path, png_path))
print('=' * 30)
cs = build_control_structure(cfg, cfg.entry_node)
cs_marks = {}
cs_str = cs_tree_to_str(cs, cs_marks)
print(cs_str)
print('=' * 30)
print_structured_flow(fn, cfg, cfg.entry_node, cs_marks)
return cs_str
except:
import traceback
traceback.print_exc()
print("Unexpected error:", sys.exc_info()[0])
print("%s had an error" % name)
return ''
def doit(fn, name=None):
if not name:
name = fn.__name__
print(name)
bb_mgr = basic_blocks(fn)
# for bb in bb_mgr.bb_list:
# print("\t", bb)
dis.dis(fn)
cfg = ControlFlowGraph(bb_mgr)
write_dot(name, "/tmp/flow-", cfg.graph, write_png=True)
try:
dt = DominatorTree(cfg)
cfg.dom_tree = dt.tree(False)
dfs_forest(cfg.dom_tree, False)
build_dom_set(cfg.dom_tree, False)
write_dot(name, "/tmp/flow-dom-", cfg.dom_tree, write_png=True)
cfg.pdom_tree = dt.tree(True)
dfs_forest(cfg.pdom_tree, True)
build_dom_set(cfg.pdom_tree, True)
write_dot(name, "/tmp/flow-pdom-", cfg.pdom_tree, write_png=True)
print("=" * 30)
augmented_instrs = augment_instructions(fn, cfg, opc.version_tuple)
for inst in augmented_instrs:
print(inst.disassemble(opc))
# return cs_str
except:
import traceback
traceback.print_exc()
print("Unexpected error:", sys.exc_info()[0])
print("%s had an error" % name)
return ""
def try_create_cfg(func, definition_line=None, args=None) -> Optional[CFG]:
try:
bb_mgr = basic_blocks(PYTHON_VERSION, IS_PYPY, func)
cfg = ControlFlowGraph(bb_mgr)
except Exception:
logger.exception("Error during CFG creation")
return None
byte_blocks_graph = cfg.graph
g = nx.MultiDiGraph()
instructions = list(_get_instructions(func, definition_line))
fake_instructions = _try_fake_instructions_function_arguments(func, definition_line=definition_line, args=args)
if fake_instructions is None:
return None
instructions.extend(fake_instructions)
offset_to_instruction = {}
for i in instructions:
offset_to_instruction[i.offset] = i
if isinstance(i, Instruction):
i.argval = 0
offset_to_line = {}
for i in instructions:
offset = i.offset
instr = i
while not instr.starts_line:
offset -= 2
instr = offset_to_instruction[offset]
offset_to_line[i.offset] = instr.starts_line
for instr in instructions:
attrs = {
LINE_KEY: offset_to_line[instr.offset],
INSTRUCTION_KEY: instr
}
g.add_node(instr.offset, **attrs)
for node in byte_blocks_graph.nodes:
start = node.bb.index[0]
end = node.bb.index[1]
source = start
for offset in range(start + 2, end + 2, 2):
for ins in instructions:
if ins.offset == offset:
g.add_edge(source, offset)
source = offset
for edge in byte_blocks_graph.edges:
remove = False
if edge.kind == 'fallthrough' \
and \
BB_JUMP_UNCONDITIONAL in edge.source.flags:
remove = True
#
# if edge.kind == 'fallthrough' \
# and \
# BB_NOFOLLOW in edge.dest.flags:
# remove = True
if edge.kind == 'no fallthrough':
remove = True
# if edge.kind == 'fallthrough':
# remove = True
# if BB_FINALLY in edge.source.flags and BB_END_FINALLY in edge.dest.flags:
# remove = True
# if edge.kind == "forward" and BB_STARTS_POP_BLOCK in edge.dest.flags:
# remove = True
if edge.kind == "exception" and BB_END_FINALLY in edge.dest.flags:
remove = True
if edge.kind == "forward" and BB_END_FINALLY in edge.dest.flags:
remove = True
if remove:
continue
def fl2name(fls):
return ", ".join([FLAG2NAME[f] for f in fls])
source = edge.source
dest = edge.dest
source_offset = source.bb.index[1]
dest_offset = dest.bb.index[0]
if not dest_offset % 2 == 0: # maybe don't add an extra node for exit?
g.add_edge(source_offset, exit_node_offset,
label=edge.kind + " / " + str(fl2name(edge.source.flags)) + " / " + str(
fl2name(edge.dest.flags)),
source_flags=fl2name(edge.source.flags),
dest_flags=fl2name(edge.dest.flags),
color="red" if remove else "black"
)
else:
# colors = ['#35332C', '#E1DAC2', '#D1B04B', '#3C2E00', '#C19200', '#242C24', '#A1B99F', '#48AB3D', '#053200',
# '#0F9E00', '#352D2C', '#E1C6C2', '#D15A4B', '#3C0700', '#C11600', '#201F25', #'#8B8B9D', '#454292',
# '#02002A', '#0C0787']
colors = ["red", "yellow", "blue", "orange", "purple", "brown"]
g.add_edge(source_offset, dest_offset,
label=edge.kind + " / " + str(fl2name(edge.source.flags)) + " / " + str(
fl2name(edge.dest.flags)),
source_flags=fl2name(edge.source.flags),
dest_flags=fl2name(edge.dest.flags),
# color="red" if remove else "black"
color=colors[hash(edge.kind) % len(colors)],
weight=10000
)
entry_node_label = 0
if len(fake_instructions) > 0:
fake_instructions = list(sorted(fake_instructions, key=lambda i: i.offset))
entry_node_label = fake_instructions[0].offset
for i1, i2 in zip(fake_instructions, fake_instructions[1:]):
g.add_edge(i1.offset, i2.offset)
last_fake_instruction_offset = fake_instructions[-1].offset
first_real_instruction_offset = 0
g.add_edge(last_fake_instruction_offset, first_real_instruction_offset)
exit_node_label = 1
function_name = func.__name__
def rename_node(node):
return function_name + "@" + str(node)
mapping = {node: rename_node(node) for node in g.nodes()}
nx.relabel_nodes(g, mapping=mapping, copy=False)
nx.set_node_attributes(g, function_name, FUNCTION_KEY)
return CFG(g, rename_node(entry_node_label), rename_node(exit_node_label))