pincher

#!/usr/bin/env python3
import argparse
import subprocess
import os
import re

filepath = os.path.dirname(os.path.realpath(__file__))

FLOATING_POINT_REGS = {"st0", "st1", "st2", "st3", "st4", "st5", "st6",
                       "st7", "xmm0", "xmm1", "xmm2", "xmm3", "xmm4", "xmm5", "xmm6", "xmm7"}


def fix_bpx(arg):
    set_regex = r"set\(([a-zA-Z0-9]+),\ *([0-9a-fA-FxX]+)\)"
    set_fp_regex = r"set\(([a-zA-Z0-9]+),\ *([0-9]+\.[0-9]+)\)"
    dmp_reg_regex = r"dump\(([a-zA-Z0-9]+),\ *([0-9a-fA-FxX]+)\)"

    set_matches = re.findall(set_regex, arg)
    set_fp_matches = re.findall(set_fp_regex, arg)
    dmp_reg_matches = re.findall(dmp_reg_regex, arg)

    if (set([x[0] for x in set_fp_matches]) - FLOATING_POINT_REGS != set()):
        not_valid_regs = set([x[0]
                              for x in set_fp_matches]) - FLOATING_POINT_REGS
        print(not_valid_regs, "are not valid floating point registers")
        exit(1)

    address = arg.split(",")[0]
    other_args = ','.join([x for x in arg.split(",") if ":" in x])

    if other_args:
        res = address + "," + other_args
    else:
        res = address
    regs = '-'.join([x[0].lower() for x in set_matches])
    vals = '-'.join([x[1].lower() for x in set_matches])

    if (regs):
        res += ",set_regs:" + regs
        res += ",set_vals:" + vals

    regs = '-'.join([x[0].lower() for x in set_fp_matches])
    vals = '-'.join([x[1].lower() for x in set_fp_matches])

    if (regs):
        res += ",set_fp_regs:" + regs
        res += ",set_fp_vals:" + vals

    regs = '-'.join([x[0].lower() for x in dmp_reg_matches])
    lens = '-'.join([x[1].lower() for x in dmp_reg_matches])

    if (regs):
        res += ",dump_regs:" + regs
        res += ",dump_lengths:" + lens

    return res


def fix_bpf(arg):
    return arg.replace("skip", "skip:1").replace("bt", "bt:1")


def check_elf(arg):
    with open(arg, "rb") as fin:
        magic = fin.read(5)

    if magic != b"\x7fELF\x02":
        print("unsupported filetype")
        exit(1)


def main(args):
    check_elf(args.prog[0])

    cmd = [os.path.join(filepath, "pin/pin"),
           "-t",
           os.path.join(filepath, "pintool/obj-intel64/tracer.so")]
    if args.bpf is not None:
        for bpf in args.bpf:
            cmd.append("-bpf")
            cmd.append(fix_bpf(bpf))

    if args.bpx is not None:
        for bpx in args.bpx:
            cmd.append("-bpx")
            cmd.append(fix_bpx(bpx))

    if args.symbs is not None:
        cmd.append("-print_symb")
        cmd.append(args.symbs)

    if args.callgraph is not None:
        cmd.append("-dump_callgraph")
        cmd.append(args.callgraph)

    cmd.append("--")
    cmd.extend(args.prog)

    subprocess.call(cmd)


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description='Reverse engineering tool for modifying/logging stuff during execution')
    parser.add_argument('--bpf', metavar='dict', type=str, nargs='*',
                        help='function breakpoint, described as key:val,... list')
    parser.add_argument('--bpx', metavar='dict', type=str, nargs='*',
                        help='instruction breakpoint, described as key:val,... list')
    parser.add_argument('--symbs', metavar='regex', type=str,
                        nargs='?', help='print symbols that match the regex')
    parser.add_argument('--callgraph', metavar='outfile.dot', type=str,
                        nargs='?', help='generate the callgraph and dump it in outfile')
    parser.add_argument('prog', type=str, nargs='+',
                        help='command line of analyzed program')
    args = parser.parse_args()

    main(args)