#!/usr/bin/env python2

import ast
import errno
import os
import subprocess
import sys
import jinja2

base_path = os.path.abspath(os.path.dirname(__file__))
for t in [os.path.join(base_path, 'templates'),
          os.path.join(sys.prefix, 'share/wuschl/templates'),
          os.path.join(sys.prefix, 'local/share/wuschl/templates')]:
    if os.path.exists(t):
        template_path = t
        break
else:
    raise RuntimeError("Cannot find wuschl's templates.")
template_loader = jinja2.FileSystemLoader(template_path)
template_env = jinja2.Environment(loader=template_loader)

def _to_bin(binary):
    return ''.join([ '\\x%02x' % ord(b) for b in binary])

class Fuzzy(object):
    def __init__(self, name):
        self.name = name
        self.afldir = name + '_afl'
        self.inputdir = self.afldir + '/input'
        for d in self.afldir, self.inputdir:
            try:
                os.mkdir(d)
            except OSError, e:
                if e.errno != errno.EEXIST:
                    raise
        self.testcases = []

    def _dict(self):
        return dict([ (key, getattr(self, key)) for key in dir(self) if not key.startswith('_')
                                                           and not callable(getattr(self, key))])

    def _render(self, suffix, template):
        with open(self.name + suffix, 'w') as ofile:
            ofile.write(template_env.get_template(template).render(self._dict()))

    def _collect_from_afl(self):
        self.testcases = []
        if not os.path.exists(self.afldir + '/queue'):
            return

        subprocess.check_call(['afl-cmin', '-i', self.afldir + '/queue', '-o', self.afldir + '/queue-min',
                               '--', './%s' % self.name, '-r'])

        for f in os.listdir(self.afldir + '/queue-min'):
            with open(os.path.join(self.afldir, 'queue-min', f), 'rb') as input_file:
                input_data = input_file.read()
            p = subprocess.Popen(['./%s' % self.name, '-r'], stdin=subprocess.PIPE,
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            output_data,_ = p.communicate(input_data)
            p.wait()
            ret = p.returncode
            self.testcases.append({
                'input': _to_bin(input_data),
                'input_len': len(input_data),
                'output': _to_bin(output_data),
                'output_len': len(output_data),
                'ret': ret
            })

    def _collect_from_prog(self, update_output):
        output = subprocess.check_output(['./%s' % self.name, '-d'])

        self.testcases = []
        for line in output.splitlines():
            line = line.strip()
            if not line:
                continue

            input_hex,output_hex,ret = line.split(',')
            self.testcases.append({
                'input': input_hex,
                'input_len': len(input_hex)//4,
                'output': output_hex,
                'output_len': len(output_hex)//4,
                'ret': int(ret)
            })

        if not update_output:
            return

        for t in self.testcases:
            input_data = ast.literal_eval('"' + t['input'] + '"')
            p = subprocess.Popen(['./%s' % self.name, '-r'], stdin=subprocess.PIPE,
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            output_data,_ = p.communicate(input_data)
            p.wait()
            ret = p.returncode

            t['output'] = _to_bin(output_data)
            t['output_len'] = len(output_data)
            t['ret'] = ret

    def create(self):
        '''Creates a new empty test from the template'''
        if os.path.exists(self.name + '.c'):
            print >>sys.stderr, "Won't override output file. Delete %s if you really want to." % (
                    self.name + '.c'
            )
            return 1
        self._render('.c', 'main.c.j2'),
        self.update()
        return 0

    def update(self):
        '''Collects tests from the afl corpus and updates the header'''
        self._collect_from_afl()
        self._render('_tests.h', 'test.h.j2')
        return 0

    def upgrade(self):
        '''Upgrades the header with the new version in wuschl. Tests stay the same'''
        self._collect_from_prog(False)
        self._render('_tests.h', 'test.h.j2')
        return 0

    def rebuild(self):
        '''Rebuilds the expected outputs for the existing tests in the header'''
        self._collect_from_prog(True)
        self._render('_tests.h', 'test.h.j2')
        return 0

    def fill_input(self):
        '''Put existing testcases into input dir'''
        self._collect_from_prog(False)
        for idx,t in enumerate(self.testcases):
            input_data = ast.literal_eval('"' + t['input'] + '"')
            with open(os.path.join(self.inputdir, 'input_%d' % idx), 'wb') as input_file:
                input_file.write(input_data)

    def fuzz(self):
        '''Starts afl for the given program'''
        if not os.listdir(self.inputdir):
            print >>sys.stderr, "Please create one or more input files and put them into %s" % self.inputdir
            return 1
        os.execvp("afl-fuzz", ["afl-fuzz", "-i", self.inputdir, "-o", self.afldir, "--",
                               "./%s" % self.name, "-r"])
        return 1

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print >>sys.stderr, "Usage: %s <op> <name>" % (sys.argv[0])
        print >>sys.stderr, "With op as one of:\n"
        for op_name in dir(Fuzzy):
            if op_name.startswith('_'):
                continue
            op = getattr(Fuzzy, op_name, None)
            if not callable(op) or not op.__doc__:
                continue
            print >>sys.stderr, "  %10s  %s" % (op_name, op.__doc__)
        print >>sys.stderr, ""
        sys.exit(1)
    f = Fuzzy(sys.argv[2])
    op = getattr(f, sys.argv[1], None)
    if not callable(op):
        print >>sys.stderr, "Unknown operation %r" % sys.argv[1]
        sys.exit(1)
    sys.exit(op())
