#!/usr/bin/env python3
#
# Interpret a file that crashes an fuzz_ndr_X binary.
#
# Copyright (C) Catalyst IT Ltd. 2019


import sys
import os
from base64 import b64encode
import struct
import argparse
import re

TYPE_MASK = 3
TYPES = ['struct', 'in', 'out']

FLAGS = [
    (4, 'ndr64', '--ndr64'),
]


def print_if_verbose(*args, **kwargs):
    if verbose:
        print(*args, **kwargs)


def process_one_file(f):
    print_if_verbose(f.name)
    print_if_verbose('-' * len(f.name))

    b = f.read()
    flags, function = struct.unpack('<HH', b[:4])
    if opnum is not None and opnum != function:
        return

    t = TYPES[flags & TYPE_MASK]
    if ndr_type and ndr_type != t:
        return

    payload = b[4:]
    data64 = b64encode(payload).decode('utf-8')

    cmd = ['bin/ndrdump',
           pipe,
           str(function),
           t,
           '--base64-input',
           '--input', data64,
    ]

    for flag, name, option in FLAGS:
        if flags & flag:
            print_if_verbose("flag: %s" % name)
            cmd.append(option)

    print_if_verbose("length: %d\n" % len(payload))
    print(' '.join(cmd))
    print_if_verbose()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--pipe', default='$PIPE',
                        help='pipe name (for output command line)')
    parser.add_argument('-t', '--type', default=None, choices=TYPES,
                        help='restrict to this type')
    parser.add_argument('-o', '--opnum', default=None, type=int,
                        help='restrict to this function/struct number')
    parser.add_argument('FILES', nargs='*', default=(),
                        help="read from these files")
    parser.add_argument('-k', '--ignore-errors', action='store_true',
                        help='do not stop on errors')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='say more')
    parser.add_argument('-H', '--honggfuzz-file',
                        help="extract crashes from this honggfuzz report")
    parser.add_argument('-f', '--crash-filter',
                        help="only print crashes matching this rexexp")

    args = parser.parse_args()

    global pipe, opnum, ndr_type, verbose
    pipe = args.pipe
    opnum = args.opnum
    ndr_type = args.type
    verbose = args.verbose

    if not args.FILES and not args.honggfuzz_file:
        parser.print_usage()
        sys.exit(1)

    for fn in args.FILES:
        if args.crash_filter is not None:
            if not re.search(args.crash_filter, fn):
                print_if_verbose(f"skipping {fn}")
                continue
        try:
            if fn == '-':
                process_one_file(sys.stdin)
            else:
                with open(fn, 'rb') as f:
                    process_one_file(f)
        except Exception:
            print_if_verbose("Error processing %s\n" % fn)
            if args.ignore_errors:
                continue
            raise

    if args.honggfuzz_file:
        print_if_verbose(f"looking at {args.honggfuzz_file}")
        with open(args.honggfuzz_file) as f:
            pipe = None
            crash = None
            for line in f:
                m = re.match(r'^\s*fuzzTarget\s*:\s*bin/fuzz_ndr_(\w+)\s*$', line)
                if m:
                    pipe = m.group(1).split('_TYPE_', 1)[0]
                    print_if_verbose(f"found pipe {pipe}")
                m = re.match(r'^FUZZ_FNAME: (\S+)$', line)
                if m:
                    crash = m.group(1)
                    if args.crash_filter is not None:
                        if not re.search(args.crash_filter, crash):
                            print_if_verbose(f"skipping {crash}")
                            pipe = None
                            crash = None
                            continue
                    print_if_verbose(f"found crash {crash}")
                if pipe is not None and crash is not None:
                    with open(crash, 'rb') as f:
                        process_one_file(f)
                    pipe = None
                    crash = None


main()