import string
import z3

from acsploit.options import Options

options = Options()
options.add_option('hash', '+ * x y z', 'The hash function to use, in prefix notation')
options.add_option('target_type', 'image', 'Whether to solve for inputs that hash to image or inputs that hash to the same value as preimage', ['image', 'preimage'])
options.add_option('image', 0, 'The target value to solve the inputs to the hash function for')
options.add_option('preimage', 'x = 4, y = 1, z = 1', 'The preimage values to hash, in the format VAR = VALUE')
options.add_option('variable_width', 32, 'Bit-width of variables in the hash function')
options.add_option('n_collisions', 10, 'The number of colliding inputs to solve for')

DESCRIPTION = 'Produces hash collisions for a custom hash function.' \
              '\n\n  ' \
              'This module allows you to define custom hash functions out of a collection of primitive operations ' \
              'and solve for collisions against them using z3. Hash functions are specified in prefix notation, using ' \
              'arbitrary variable names and the following operators: +, -, *, /, <<, >>, & (bit-wise AND), | (bit-wise OR), ' \
              'and ^ (bit-wise XOR). See custom_hashes.md for more details.'

NO_INPUT = True


def run(output):
    ast = parse_input(options['hash'])
    variables = {}  # map from names to z3_vars
    z3_expression = ast.convert_to_z3(variables)

    solver = z3.Solver()
    if options['target_type'] == 'image':
        solver.add(options['image'] == z3_expression)
    elif options['target_type'] == 'preimage':
        # extract and validate the user-provided preimage
        preimage = options['preimage']
        var_defs = preimage.split(',')
        variable_values = {}
        if len(var_defs) < len(variables):
            raise ValueError('Not enough preimage values given for all variables used in the equation')
        for var_def in var_defs:
            try:
                variable_name, value = var_def.split('=', 1)
            except ValueError:
                raise ValueError('Invalid syntax for preimage values')
            variable_name = variable_name.strip()
            if variable_name in variable_values:
                raise ValueError('Multiple preimage values given for variable "%s"' % variable_name)
            try:
                value = int(value)
            except ValueError:
                raise ValueError('Preimage value "%s" for variable "%s" is not an integer' % (value, variable_name))
            variable_values[variable_name] = value
        for variable_name in variables:
            if variable_name not in variable_values:
                raise ValueError('Preimage value not given for variable "%s"' % variable_name)

        # we have a preimage but we want an image to set z3_expression equal to for solving
        # so we set a new variable equal to z3_expression, provide the preimage values,
        #  and then ask Z3 to solve for our new variable
        target_var = z3.BitVec('__v', ast.target_width)
        sub_solver = z3.Solver()
        sub_solver.add(target_var == z3_expression)
        for variable in variables:
            sub_solver.add(variables[variable] == variable_values[variable])
        assert sub_solver.check() == z3.sat  # this should always hold, since the target var is unbounded
        solution = sub_solver.model()
        target_value = solution[target_var]

        # we can now set z3_expression equal to the appropriate image
        solver.add(target_value == z3_expression)
        # and also prevent the preimage values being generated as a solution
        solver.add(z3.Or([var() != solution[var] for var in solution if var.name() != '__v']))

    solutions = []
    while solver.check() == z3.sat and len(solutions) < options['n_collisions']:
        solution = solver.model()
        solutions.append(solution)
        # prevent duplicate solutions
        solver.add(z3.Or([var() != solution[var] for var in solution]))

    output.output(solutions)


class Node(object):
    CONSTANT = 0
    VARIABLE = 1

    ADDITION = 10
    SUBTRACTION = 11
    MULTIPLICATION = 12
    DIVISION = 13

    # arithmetic shifts
    LEFT_SHIFT = 20
    RIGHT_SHIFT = 21

    # bitwise logic
    AND = 30
    OR = 31
    XOR = 32

    VARIABLE_CHARS = string.ascii_lowercase
    OPERATION_CHARS = {
        '+': ADDITION,
        '-': SUBTRACTION,
        '*': MULTIPLICATION,
        '/': DIVISION,
        '<<': LEFT_SHIFT,
        '>>': RIGHT_SHIFT,
        '&': AND,
        '|': OR,
        '^': XOR
    }

    def __init__(self, operation, left, right=None):
        self.operation = operation
        self.left = left
        self.right = right
        self.variable_width = options['variable_width']  # bit-width of variables
        self.target_width = self.variable_width  # bit-width of calculations
        if self.variable_width <= 32:  # if under 32-/64-bit, assume user wants one of those
            self.target_width = 32
        elif self.variable_width <= 64:
            self.target_width = 64

    def __str__(self):
        s = 'Node: %i ' % id(self)
        s += '(Op: %s, ' % str(self.operation)
        if self.operation in [Node.CONSTANT, Node.VARIABLE]:
            s += 'Value: %s)' % str(self.left)
            assert self.right is None  # TODO: this should be elsewhere
        else:  # actual operation
            s += 'Left: %i, ' % id(self.left)
            s += 'Right: %i)' % id(self.right)
        return s

    def print_tree(self):
        print(str(self))
        if type(self.left) is Node:
            self.left.print_tree()
        if type(self.right) is Node:
            self.right.print_tree()

    def convert_to_z3(self, variables):
        if self.operation == Node.CONSTANT:
            return z3.BitVecVal(self.left, self.target_width)
        elif self.operation == Node.VARIABLE:
            if self.left in variables:
                return variables[self.left]
            else:
                var = z3.BitVec(self.left, self.variable_width)
                ext_var = z3.ZeroExt(self.target_width - self.variable_width, var)
                variables[self.left] = ext_var
                return ext_var
        else:  # operation
            left = self.left.convert_to_z3(variables)
            right = self.right.convert_to_z3(variables)
            if self.operation == Node.ADDITION:
                return left + right
            elif self.operation == Node.SUBTRACTION:
                return left - right
            elif self.operation == Node.MULTIPLICATION:
                return left * right
            elif self.operation == Node.DIVISION:
                return left / right
            elif self.operation == Node.LEFT_SHIFT:
                return left << right
            elif self.operation == Node.RIGHT_SHIFT:
                return left >> right
            elif self.operation == Node.AND:
                return left & right
            elif self.operation == Node.OR:
                return left | right
            elif self.operation == Node.XOR:
                return left ^ right
        raise LookupError('Unknown operation: %s' % str(self.operation))

    @staticmethod
    def make_constant(value):
        return Node(Node.CONSTANT, value)

    @staticmethod
    def make_variable(variable_name):
        return Node(Node.VARIABLE, variable_name)

    @staticmethod
    def make_operation(operation, left, right):
        return Node(operation, left, right)


def parse_input(input):
    atoms = input.split(' ')
    tree, atoms = parse_atom(atoms)
    assert len(atoms) == 0
    return tree


def parse_atom(atoms):
    assert len(atoms) > 0

    constant = parse_constant(atoms[0])
    if constant is not None:
        return Node.make_constant(constant), atoms[1:]

    variable = parse_variable(atoms[0])
    if variable is not None:
        return Node.make_variable(variable), atoms[1:]

    operation = parse_operation(atoms[0])
    if operation is not None:
        left, atoms = parse_atom(atoms[1:])  # pop operation off of atoms
        right, atoms = parse_atom(atoms)
        return Node.make_operation(operation, left, right), atoms


def parse_constant(atom):
    # only support base 10 for now
    try:
        return int(atom)
    except ValueError:
        return None


def parse_variable(atom):
    for char in atom:
        if char not in Node.VARIABLE_CHARS:
            return None
    return atom


def parse_operation(atom):
    try:
        return Node.OPERATION_CHARS[atom]
    except KeyError:
        return None
