#!/usr/bin/python
# -*- coding: utf-8 -*-


import argparse
import sys

# from pychrysalide.features import *
from pychrysalide.analysis.contents import FileContent
from pychrysalide.analysis import StudyProject
from pychrysalide.arch import ArchInstruction
from pychrysalide.core import wait_for_all_global_works


def init_fake_ranks(blocks):
    """Build a list of ranks for basic blocks."""

    ranks = {}

    for b in blocks:
        ranks[b] = -1

    return ranks


def dump_stack(stack):
    """Print the stack of looping code blocks."""

    for blk in stack:

        print(' Loop detected:', blk.boundaries[0].range.addr)

    sys.exit(2)


def track_infinite_loop_when_ranking(blk, ranks, stack):
    """Mimic the rank_routine_block() C function."""

    fatal = blk in stack

    stack.append(blk)

    if fatal:
        dump_stack(stack)

    next_rank = ranks[blk] + 1

    for dest, dtype in blk.destinations:

        if dtype == ArchInstruction.ILT_LOOP:
            continue

        if dtype != ArchInstruction.ILT_EXEC_FLOW \
           and dtype != ArchInstruction.ILT_JUMP \
           and dtype != ArchInstruction.ILT_CASE_JUMP \
           and dtype != ArchInstruction.ILT_JUMP_IF_TRUE \
           and dtype != ArchInstruction.ILT_JUMP_IF_FALSE:
            continue

        if next_rank > ranks[blk] or ranks[blk] == -1:

            ranks[dest] = next_rank

            track_infinite_loop_when_ranking(dest, ranks, stack)

    stack.remove(blk)


if __name__ == '__main__':

    title = '%s - Track infinite loops when computing ranks.' % sys.argv[0]
    title += '\n\n'
    title += '(this is a debug helper tool: call to rank_routine_blocks() should be disabled in src/analysis/disass/routines.c)'

    parser = argparse.ArgumentParser(description=title, add_help=False, formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('-h', '--help', action='store_true', help='Display the command line options understood by %s.' % sys.argv[0])

    parser.add_argument('binfile', type=str, help='The object file to be examined')
    parser.add_argument('fname', type=str, help='The analyzed function to process')

    args = parser.parse_args()

    if args.help:
        parser.print_help()
        sys.exit(1)

    prj = StudyProject()

    cnt = FileContent(args.binfile)

    prj.discover(cnt)

    wait_for_all_global_works()

    binary = prj.contents[0]

    sym = binary.format.find_symbol_by_label(args.fname)

    if not(sym):
        print('Function "%s" not found!' % args.fname)
        sys.exit(1)

    blocks = list(sym.basic_blocks)

    ranks = init_fake_ranks(blocks)

    ranks[blocks[0]] = 0

    track_infinite_loop_when_ranking(blocks[0], ranks, [])