#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Commands for setting temporary breakpoints on the next
instruction of some type (call, branch, etc.)
"""

import re

import capstone
import gdb

import pwndbg.disasm
import pwndbg.regs
from pwndbg.color import message

jumps = set((
    capstone.CS_GRP_CALL,
    capstone.CS_GRP_JUMP,
    capstone.CS_GRP_RET,
    capstone.CS_GRP_IRET
))

interrupts = set((capstone.CS_GRP_INT,))

@pwndbg.events.exit
def clear_temp_breaks():
    if not pwndbg.proc.alive:
        breakpoints = gdb.breakpoints()
        if breakpoints:
            for bp in breakpoints:
                if bp.temporary and not bp.visible: #visible is used instead of internal because older gdb's don't support internal 
                    bp.delete()

def next_int(address=None):
    """
    If there is a syscall in the current basic black,
    return the instruction of the one closest to $PC.

    Otherwise, return None.
    """
    if address is None:
        ins = pwndbg.disasm.one(pwndbg.regs.pc)
        if not ins:
            return None
        address = ins.next

    ins = pwndbg.disasm.one(address)
    while ins:
        if set(ins.groups) & jumps:
            return None
        if set(ins.groups) & interrupts:
            return ins
        ins = pwndbg.disasm.one(ins.next)

    return None


def next_branch(address=None):
    if address is None:
        ins = pwndbg.disasm.one(pwndbg.regs.pc)
        if not ins:
            return None
        address = ins.next

    ins = pwndbg.disasm.one(address)
    while ins:
        if set(ins.groups) & jumps:
            return ins
        ins = pwndbg.disasm.one(ins.next)

    return None


def break_next_branch(address=None):
    ins = next_branch(address)

    if ins:
        gdb.Breakpoint("*%#x" % ins.address, internal=True, temporary=True)
        gdb.execute('continue', from_tty=False, to_string=True)
        return ins


def break_next_interrupt(address=None):
    ins = next_int(address)

    if ins:
        gdb.Breakpoint("*%#x" % ins.address, internal=True, temporary=True)
        gdb.execute('continue', from_tty=False, to_string=True)
        return ins


def break_next_call(symbol_regex=None):
    while pwndbg.proc.alive:
        ins = break_next_branch()

        if not ins:
            break

        # continue if not a call
        if capstone.CS_GRP_CALL not in ins.groups:
            continue

        # return call if we don't search for a symbol
        if not symbol_regex:
            return ins

        # return call if we match target address
        if ins.target_const and re.match('%s$' % symbol_regex, hex(ins.target)):
            return ins

        # return call if we match symbol name
        if ins.symbol and re.match('%s$' % symbol_regex, ins.symbol):
            return ins


def break_next_ret(address=None):
    while pwndbg.proc.alive:
        ins = break_next_branch(address)

        if not ins:
            break

        if capstone.CS_GRP_RET in ins.groups:
            return ins


def break_on_program_code():
    """
    Breaks on next instruction that belongs to process' objfile code.
    :return: True for success, False when process ended or when pc is at the code.
    """
    mp = pwndbg.proc.mem_page
    start = mp.start
    end = mp.end

    if start <= pwndbg.regs.pc < end:
        print(message.error('The pc is already at the binary objfile code. Not stepping.'))
        return False

    while pwndbg.proc.alive:
        gdb.execute('si', from_tty=False, to_string=False)

        addr = pwndbg.regs.pc
        if start <= addr < end:
            return True

    return False


def break_on_next(address=None):
    address = address or pwndbg.regs.pc
    ins = pwndbg.disasm.one(address)

    gdb.Breakpoint("*%#x" % (ins.address + ins.size), temporary=True)
    gdb.execute('continue', from_tty=False, to_string=True)
