Python z3.BitVec() Examples

def find_name_for_bit_at_index(index, bit):
    solver = z3.Solver()

    NAME_LENGTH = 12  # arbitrary
    name = [z3.BitVec("c%d" % i, 32) for i in range(NAME_LENGTH)]
    for i in range(len(name)):
        solver.add(z3.And(name[i] > 0, name[i] <= 0xff))

    h1 = hash1(name)
    solver.add(h1 == index)

    h2 = hash2(name)
    solver.add(h2 == bit)

    h3 = hash3(name)
    solver.add(z3.Extract(15, 0, h3) == h2)  # for simplicity

    if solver.check() == z3.sat:
        return "".join(chr(solver.model()[c].as_long()) for c in name).encode("latin-1") 

def get_sym_bitvec(self, constraint_type, gen, bv_size=256, unique=False, **kwargs):
        vector_name = ConstraintType[constraint_type.name].value
        label_template = vector_name + '_gen{}'
        for k in kwargs:
            label_template += '_' + k + '{' + k + '}'
        label = label_template.format(gen, **kwargs)
        if unique:
            unique_id = self.unique_vector_name_counter.get(vector_name, 0)
            self.unique_vector_name_counter[vector_name] = unique_id + 1
            label = label + '_uid' + str(unique_id)
        assert constraint_type != ConstraintType.CALLDATA or 'acc' not in kwargs

        if constraint_type == ConstraintType.CALLDATA_ARRAY:
            return z3.Array(label, z3.BitVecSort(bv_size), z3.BitVecSort(8))
        elif constraint_type in [ConstraintType.CALLER, ConstraintType.ORIGIN, ConstraintType.ENTRY_ACCOUNT]:
            return svm_utils.zpad_bv_right(z3.BitVec(label, svm_utils.ADDRESS_LEN), svm_utils.VECTOR_LEN)
        else:
            return z3.BitVec(label, bv_size) 

def to_z3_variable(self, v):
        """
        Transforms a variable into a z3 variable
        :param v: variable
        :return: z3 variable
        """
        # no variable
        if v not in self.variables:
            return v

        # get variable
        v = self.variables[v]
        # initialise dict
        if v.name not in self._z3_var:
            v_z3 = z3.BitVec(v.name, v.size)
            self._z3_var[v.name] = v_z3

        # get z3 variable
        v_z3 = self._z3_var[v.name]

        return v_z3 

def __init__(self,
                 address,
                 contract,
                 typ=AccountType.DEFAULT,
                 balance=None,
                 storage=None,
                 account_id=None,
                 mapping_id_to_sum=None):
        global account_counter
        if account_id is None:
            account_counter += 1
        self.id = account_id if account_id is not None else account_counter
        self.address = address
        self.contract = contract
        self.typ = typ
        self.balance = z3.BitVec(f'{address}_balance', 256) if balance is None else balance
        self.CONTRACT_TO_ACCOUNT_COUNT[contract] += 1
        self.contract_tag = contract.name + utils.ADDRESS_ARG_TAG + str(self.CONTRACT_TO_ACCOUNT_COUNT[contract])
        self.storage = storage if storage is not None else EmptyStorage()
        self.mapping_id_to_sum = mapping_id_to_sum if mapping_id_to_sum is not None else {} 

def get(self, addr, size):
        """ Memory access at address @addr of size @size.
        @addr: a z3 BitVec, the address to read.
        @size: int, size of the read in bits.
        Return a z3 BitVec of size @size representing a memory access.
        """
        original_size = size
        if original_size % 8 != 0:
            # Size not aligned on 8bits -> read more than size and extract after
            size = ((original_size // 8) + 1) * 8
        res = self[addr]
        if self.is_little_endian():
            for i in range(1, size // 8):
                res = z3.Concat(self[addr+i], res)
        else:
            for i in range(1, size //8):
                res = z3.Concat(res, self[addr+i])
        if size == original_size:
            return res
        else:
            # Size not aligned, extract right sized result
            return z3.Extract(original_size-1, 0, res) 

def update_eax(self):
        self.set_reg('eax', z3.BitVec('ret', 32)) 

def check_eq_z3(self, target, pattern):
        'Check equivalence with z3'
        # pylint: disable=exec-used
        getid = asttools.GetIdentifiers()
        getid.visit(target)
        if getid.functions:
            # not checking exprs with functions for now, because Z3
            # does not seem to support function declaration with
            # arbitrary number of arguments
            return False
        for var in self.variables:
            exec("%s = z3.BitVec('%s', %d)" % (var, var, self.nbits))
        target_ast = deepcopy(target)
        target_ast = Unflattening().visit(target_ast)
        ast.fix_missing_locations(target_ast)
        code1 = compile(ast.Expression(target_ast), '<string>', mode='eval')
        eval_pattern = deepcopy(pattern)
        EvalPattern(self.wildcards).visit(eval_pattern)
        eval_pattern = Unflattening().visit(eval_pattern)
        ast.fix_missing_locations(eval_pattern)
        getid.reset()
        getid.visit(eval_pattern)
        if getid.functions:
            # same reason as before, not using Z3 if there are
            # functions
            return False
        gvar = asttools.GetIdentifiers()
        gvar.visit(eval_pattern)
        if any(var.isupper() for var in gvar.variables):
            # do not check if all patterns have not been replaced
            return False
        code2 = compile(ast.Expression(eval_pattern), '<string>', mode='eval')
        sol = z3.Solver()
        if isinstance(eval(code1), int) and eval(code1) == 0:
            # cases where target == 0 are too permissive
            return False
        sol.add(eval(code1) != eval(code2))
        return sol.check().r == -1 

def evaluate_expr_z3(self, expr, constraints, output_size):
        """
        Tramsforms an expression to z3
        :param expr: str
        :param constraints: list of constraints
        :return: int
        """
        # to z3 expression
        expr = self.to_z3(expr)
        # initialise solver
        solver = z3.Solver()
        # output variable
        output = z3.BitVec("o", output_size)

        solver.add(expr == output)

        # add constraints
        for c in constraints:
            solver.add(c)

        # check sat
        assert (solver.check() == z3.sat)
        # parse output
        ret = solver.model()[output].as_long()

        return ret 

def get_model(self, target, pattern):
        'When target is constant and wildcards have no value yet'
        # pylint: disable=exec-used
        if target.n == 0:
            # zero is too permissive
            return False
        getwild = asttools.GetIdentifiers()
        getwild.visit(pattern)
        if getwild.functions:
            # not getting model for expr with functions
            return False
        wilds = getwild.variables
        # let's reduce the model to one wildcard for now
        # otherwise it adds a lot of checks...
        if len(wilds) > 1:
            return False

        wil = wilds.pop()
        if wil in self.wildcards:
            if not isinstance(self.wildcards[wil], ast.Num):
                return False
            folded = deepcopy(pattern)
            folded = Unflattening().visit(folded)
            EvalPattern(self.wildcards).visit(folded)
            folded = asttools.ConstFolding(folded, self.nbits).visit(folded)
            return folded.n == target.n
        else:
            exec("%s = z3.BitVec('%s', %d)" % (wil, wil, self.nbits))
        eval_pattern = deepcopy(pattern)
        eval_pattern = Unflattening().visit(eval_pattern)
        ast.fix_missing_locations(eval_pattern)
        code = compile(ast.Expression(eval_pattern), '<string>', mode='eval')
        sol = z3.Solver()
        sol.add(target.n == eval(code))
        if sol.check().r == 1:
            model = sol.model()
            for inst in model.decls():
                self.wildcards[str(inst)] = ast.Num(int(model[inst].as_long()))
            return True
        return False 

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)) 

def _generate_ascii_printable_string(base_name, size, solver):
    # establishes z3 variable names for bytes of the input string
    bytes = [z3.BitVec('%s%d' % (base_name, i), 8) for i in range(size)]
    # adds the constraint that the bytes are printable ascii
    solver.add(z3.And([_ascii_printable(byte) for byte in bytes]))
    return bytes 

def substitute_bvset(bvset, prefix):
    index = 0
    substitution_dict = {}
    for bv in bvset:
        if type(bv) == z3.BitVecRef:
            substitution_dict[bv] = z3.BitVec('{}_{}_{}'.format(bv, prefix, index), bv.size())
        elif type(bv) == z3.ArrayRef:
            substitution_dict[bv] = z3.Array('{}_{}_{}'.format(bv, prefix, index), bv.sort().domain(), bv.sort().range())
        index += 1

    return substitution_dict


# Wrapper for allowing Z3 ASTs to be stored into Python Hashtables. 

def abstract(self, label_suffix):
        old_storage = self.storage
        self.storage = AbstractStorage(f'abstract_storage{label_suffix}')
        for map_id in self.mapping_id_to_sum:
            if svm_utils.is_bv_concrete(map_id):
                map_id_string = svm_utils.get_concrete_int(map_id)
            else:
                map_id_string = str(z3.simplify(map_id))
                raise Exception("pdb")
            label = f'abstract_sum_{map_id_string}{label_suffix}'
            self.mapping_id_to_sum[map_id] = z3.BitVec(label, 256)
        self.balance = z3.BitVec(f'gstate_balance{label_suffix}', 256) 

def __new_unk_var(self):
        x = z3.BitVec('unk{0}'.format(self.__unk_var_cnt), 32)
        self.__unk_var_cnt += 1
        return x 

def CALLDATALOAD(self, gstate, index):
        logging.debug('CALLDATA index:' + str(index))
        if gstate.environment.calldata is None:
            raise SVMRuntimeError('CALLDATA is not set')
        index_concrete = svm_utils.is_bv_concrete(index)
        if gstate.environment.calldata_type == CalldataType.UNDEFINED:
            data_bytes = []
            # for i in range(32):
                # label = 'calldata_{}_{}'.format(wstate.gen, i + svm_utils.get_concrete_int(index))
                # data_bytes.append(z3.BitVec(label, 8))
            # data = z3.Concat(data_bytes)
            if index_concrete:
                data = self.svm.sym_bv_generator.get_sym_bitvec(constraints.ConstraintType.CALLDATA,
                                                            gstate.wstate.gen,
                                                            index=svm_utils.get_concrete_int(index))
                gstate.mstate.stack.append(data)
            else:
                data = self.svm.sym_bv_generator.get_sym_bitvec(constraints.ConstraintType.CALLDATA,
                                                            gstate.wstate.gen,
                                                            unique=True)
                gstate.mstate.stack.append(data)

        elif gstate.environment.calldata_type == CalldataType.DEFINED:
            assert gstate.environment.calldata.sort().name() == 'bv', 'CALLDATA sort mismatch'
            index = svm_utils.get_concrete_int(z3.simplify(index))
            calldatasize = gstate.environment.calldata.size()
            assert index < calldatasize
            calldata_start = calldatasize-index*8
            calldata_stop = max(0, calldata_start - 256)
            calldata = z3.Extract(calldata_start-1, calldata_stop, gstate.environment.calldata)
            calldata = svm_utils.zpad_bv_left(calldata, 256)
            gstate.mstate.stack.append(calldata)
        else:
            raise SVMRuntimeError('Unknown calldata type') 

def get_sym_balance(self, account):
            return z3.BitVec(ConstraintType.BALANCE.value+'_'+str(account.id), 256) 

def test_xor36(self):
        'Test that CSE of the xor36 function is equivalent to original'
        # pylint: disable=exec-used
        pwd = os.path.dirname(os.path.realpath(__file__))
        input_file = open(os.path.join(pwd, 'xor36_flat'), 'r')
        input_string = input_file.read()
        input_ast = ast.parse(input_string)
        coderef = compile(ast.Expression(input_ast.body[0].value),
                          '<string>', 'eval')
        jack = asttools.GetIdentifiers()
        jack.visit(input_ast)

        cse_string = cse.apply_cse(input_string)[0]
        # get all assignment in one ast
        assigns = cse_string[:cse_string.rfind('\n')]
        cse_assign_ast = ast.parse(assigns, mode='exec')
        assign_code = compile(cse_assign_ast, '<string>', mode='exec')
        # get final expression in one ast
        result_string = cse_string.splitlines()[-1]
        result_ast = ast.Expression(ast.parse(result_string).body[0].value)
        result_code = compile(result_ast, '<string>', mode='eval')

        for var in list(jack.variables):
            exec("%s = z3.BitVec('%s', 8)" % (var, var))
        exec(assign_code)
        sol = z3.Solver()
        sol.add(eval(coderef) != eval(result_code))
        self.assertEqual(sol.check().r, -1) 

def from_ExprLoc(self, expr):
        if self.loc_db is None:
            # No loc_db, fallback to default name
            return z3.BitVec(str(expr), expr.size)
        loc_key = expr.loc_key
        offset = self.loc_db.get_location_offset(loc_key)
        if offset is not None:
            return z3.BitVecVal(offset, expr.size)
        # fallback to default name
        return z3.BitVec(str(loc_key), expr.size) 

def from_ExprId(self, expr):
        return z3.BitVec(str(expr), expr.size) 

def arg_or_var(self, addr):
        name = get_name_from_addr_str(addr)
        if name:
            return z3.BitVec(name, 32)
        return self.__new_mem_var() 

def __getitem__(self, addr):
        """One byte memory access. Different address sizes with the same value
        will result in different memory accesses.
        @addr: a z3 BitVec, the address to read.
        Return a z3 BitVec of size 8 bits representing a memory access.
        """
        size = addr.size()
        mem = self.get_mem_array(size)
        return mem[addr] 

def __new_mem_var(self):
        x = z3.BitVec('mem{0}'.format(self.__mem_var_cnt), 32)
        self.__mem_var_cnt += 1
        return x 

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)