Skip to content

Instantly share code, notes, and snippets.

Created March 4, 2021 20:28
Show Gist options
  • Save mrphrazer/da32217f231e1dd842986f94aa6d9d37 to your computer and use it in GitHub Desktop.
Save mrphrazer/da32217f231e1dd842986f94aa6d9d37 to your computer and use it in GitHub Desktop.
Flattening Heuristic Implementation
# (c) Tim Blazytko 2021
# implementation based on the blog post "Automated Detection of Control-flow Flattening"
import sys
from miasm.analysis.binary import Container
from miasm.analysis.machine import Machine
from miasm.core.locationdb import LocationDB
def calc_flattening_score(asm_graph, loc_db):
# init score
score = 0.0
# walk over all entry nodes in the graph
for head in asm_graph.heads_iter():
# compute dominator tree
dominator_tree = asm_graph.compute_dominator_tree(head)
# walk over all basic blocks
for block in asm_graph.blocks:
# get location key for basic block via basic block address
block_key = loc_db.get_offset_location(block.lines[0].offset)
# get all blocks that are dominated by the current block
dominated = set(
[block_key] + [b for b in dominator_tree.walk_depth_first_forward(block_key)])
# check for a back edge
if not any([b in dominated for b in asm_graph.predecessors(block_key)]):
# calculate relation of dominated blocks to the blocks in the graph
score = max(score, len(dominated)/len(asm_graph.nodes()))
return score
# check args
if len(sys.argv) < 3:
print("[x] Syntax: {} <file> <addr>".format(sys.argv[0]))
# parse stdin
file_path = sys.argv[1]
start_addr = int(sys.argv[2], 16)
# init symbol table
loc_db = LocationDB()
# open the binary for analysis
container = Container.from_stream(open(file_path, 'rb'), loc_db)
# cpu abstraction
machine = Machine(container.arch)
# init disassemble engine
mdis = machine.dis_engine(container.bin_stream, loc_db=loc_db)
# disassemble the function at address
asm_cfg = mdis.dis_multiblock(start_addr)
flattening_score = calc_flattening_score(asm_cfg, loc_db)
print(f"flattening score {flattening_score} for function {hex(start_addr)}")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment