state.py

from __future__ import annotations

from insn import Instruction
from collections import deque, defaultdict
from typing import Optional
from solver import Solver
from copy import copy
from errors import *
import ops
import z3

CONST_ARR = [1, 274877906974, 244813135872, 0, 244813135872, 206158430208, 1, 90194313216, 244813135872, 1, 81604378624, 4294967295, 73014444032, 1, 68719476736, 4294327160, 81604378624, 73014444032, 244813135872, 1, 68719476736, 206158430208, 8, 94489280512, 249108103168, 283467841556, 244813135872, 240518168576, 244813135872, 287762808832, 0, 8589934598, 73014444032, 292057776130, 8589934599, 73014444032, 292057776130, 8589934600, 73014444032, 292057776130, 8589934601, 73014444032, 292057776130, 2209421562, 98784247808, 279172874242, 1095216660480, 0, 8589934602, 73014444032, 292057776130, 8589934603, 73014444032, 292057776130, 8589934604, 73014444032, 292057776130, 8589934605, 73014444032, 292057776130, 4020009855, 98784247808, 279172874242, 1095216660480, 0, 8589934606, 73014444032, 292057776130, 8589934607, 73014444032, 292057776130, 8589934608, 73014444032, 292057776130, 8589934609, 73014444032, 292057776130, 2511570847, 98784247808, 279172874242, 1095216660480, 0, 8589934610, 73014444032, 292057776130, 8589934611, 73014444032, 292057776130, 8589934612, 73014444032, 292057776130, 8589934613, 73014444032, 292057776130, 825727845, 98784247808, 279172874242, 1095216660480, 0, 8589934614, 73014444032, 292057776130, 8589934615, 73014444032, 292057776130, 8589934616, 73014444032, 292057776130, 8589934617, 73014444032, 292057776130, 2747945899, 98784247808, 279172874242, 1095216660480, 0, 8589934618, 73014444032, 292057776130, 8589934619, 73014444032, 292057776130, 8589934620, 73014444032, 292057776130, 8589934621, 73014444032, 292057776130,
             2434240953, 98784247808, 279172874242, 1095216660480, 0, 8589934622, 73014444032, 292057776130, 8589934623, 73014444032, 292057776130, 8589934624, 73014444032, 292057776130, 8589934625, 73014444032, 292057776130, 3923412385, 98784247808, 279172874242, 1095216660480, 0, 8589934626, 73014444032, 292057776130, 8589934627, 73014444032, 292057776130, 8589934628, 73014444032, 292057776130, 8589934629, 73014444032, 292057776130, 1510700589, 98784247808, 279172874242, 1095216660480, 0, 8589934630, 73014444032, 292057776130, 8589934631, 73014444032, 292057776130, 8589934632, 73014444032, 292057776130, 8589934633, 73014444032, 292057776130, 3658116609, 98784247808, 279172874242, 1095216660480, 0, 8589934634, 73014444032, 292057776130, 8589934635, 73014444032, 292057776130, 8589934636, 73014444032, 292057776130, 8589934637, 73014444032, 292057776130, 1210550661, 98784247808, 279172874242, 1095216660480, 0, 8589934638, 73014444032, 292057776130, 8589934639, 73014444032, 292057776130, 8589934640, 73014444032, 292057776130, 8589934641, 73014444032, 292057776130, 2892531646, 98784247808, 279172874242, 1095216660480, 0, 8589934642, 73014444032, 292057776130, 8589934643, 73014444032, 292057776130, 8589934644, 73014444032, 292057776130, 8589934645, 73014444032, 292057776130, 648401340, 98784247808, 279172874242, 1095216660480, 0, 8589934646, 73014444032, 292057776130, 8589934647, 73014444032, 292057776130, 8589934648, 73014444032, 292057776130, 8589934649, 73014444032, 292057776130, 2537403886, 98784247808, 279172874242, 1095216660480, 296352743424]


class State:
    path = []

    def __init__(self,
                 instructions: list[Instruction],
                 input: list
                 ):
        self.ip = 0
        self.stack = deque()
        self.locals = defaultdict()
        self.instructions = instructions
        self.solver = Solver()
        if isinstance(input[0], z3.ExprRef):
            for item in input:
                self.solver.constraints.append(item < 0xff)
                self.solver.constraints.append(item >= 0x20)
        self.operations = {
            ops.Add: self._add,
            ops.Sub: self._sub,
            ops.BitOr: self._bit_or,
            ops.BitAnd: self._bit_and,
            ops.Xor: self._xor,
            ops.Shl: self._shl,
            ops.Shr: self._shr,

            ops.Lt: self._lt,
            ops.Neq: self._neq,
            ops.Eq: self._eq,

            ops.Branch: self._branch,
            ops.BrFalse: self._br_false,

            ops.StLoc: self._st_loc,
            ops.CopyLoc: self._copy_loc,
            ops.MoveLoc: self._move_loc,
            ops.MutBorrowLoc: self._mut_borrow_loc,
            ops.ImmBorrowLoc: self._imm_borrow_loc,

            ops.LdConst: self._ld_const,
            ops.ReadRef: self._read_ref,

            ops.Pop: self._pop,

            ops.Abort: self._abort,
            ops.Ret: self._ret,

            ops.VecImmBorrow: self._vec_imm_borrow,
            ops.VecLen: self._vec_len,
            ops.VecPopBack: self._vec_pop_back,
            ops.VecPushBack: self._vec_push_back,
            ops.VecPack: self._vec_pack,

            ops.CastU8: self._cast_u8,
            ops.CastU64: self._cast_u64,
            ops.CastU128: self._cast_u128,

            ops.LdU8: self._ld_u8,
            ops.LdU64: self._ld_u64,
            ops.LdU128: self._ld_u128,
        }
        self.locals[0] = input
        self.input = input

    def __repr__(self):
        return f"State(ip={self.ip}, instructions={self.instruction})"

    @property
    def instruction(self) -> Optional[Instruction]:
        try:
            return self.instructions[self.ip]
        except KeyError:
            return None

    def clone(self) -> State:
        state = State(self.instructions, copy(self.input))
        state.ip = self.ip
        state.stack = self.stack.copy()
        state.locals = self.locals.copy()
        state.solver = self.solver.clone()
        return state

    def is_satisfiable(self) -> bool:
        return self.solver.is_satisfiable()

    def step(self) -> list[State]:
        if self.instruction is None:
            return []
        op = self.instruction.op
        arg = self.instruction.arg
        # print(f"{self.ip}: {self.instruction}")
        return self.operations[op](arg)

    def concretize(self):
        output = ""
        for var in self.input:
            value = self.solver.eval(var)
            output += chr(value)
        return output

    @classmethod
    def create_entry_state(cls, instructions: list[Instruction], input_: Optional[list] = None) -> State:
        if input_ is None:
            input_ = [z3.BitVec(f'a{x}', 64) for x in range(58)]
        return cls(instructions, input_)

    def _stack_peek(self):
        try:
            top = self.stack[-1]
        except IndexError as e:
            raise EmptyStackError() from e
        return top

    def _add(self, arg: int) -> list[State]:
        self.stack.append(self.stack.pop() + self.stack.pop())
        self.ip += 1
        return [self]

    def _sub(self, arg: int) -> list[State]:
        rhs = self.stack.pop()
        lhs = self.stack.pop()
        self.stack.append(lhs - rhs)
        self.ip += 1
        return [self]

    def _bit_or(self, arg: int) -> list[State]:
        rhs = self.stack.pop()
        lhs = self.stack.pop()
        self.stack.append(lhs | rhs)
        self.ip += 1
        return [self]

    def _bit_and(self, arg: int) -> list[State]:
        rhs = self.stack.pop()
        lhs = self.stack.pop()
        self.stack.append(lhs & rhs)
        self.ip += 1
        return [self]

    def _xor(self, arg: int) -> list[State]:
        rhs = self.stack.pop()
        lhs = self.stack.pop()
        self.stack.append(lhs ^ rhs)
        self.ip += 1
        return [self]

    def _shl(self, arg: int) -> list[State]:
        rhs = self.stack.pop()
        lhs = self.stack.pop()
        self.stack.append(lhs << rhs)
        self.ip += 1
        return [self]

    def _shr(self, arg: int) -> list[State]:
        rhs = self.stack.pop()
        lhs = self.stack.pop()
        self.stack.append(lhs >> rhs)
        self.ip += 1
        return [self]

    def _lt(self, arg: int) -> list[State]:
        rhs = self.stack.pop()
        lhs = self.stack.pop()
        self.stack.append(lhs < rhs)
        self.ip += 1
        return [self]

    def _neq(self, arg: int) -> list[State]:
        self.stack.append(self.stack.pop() != self.stack.pop())
        self.ip += 1
        return [self]

    def _eq(self, arg: int) -> list[State]:
        self.stack.append(self.stack.pop() == self.stack.pop())
        self.ip += 1
        return [self]

    def _branch(self, arg: int) -> list[State]:
        self.ip = arg
        return [self]

    def _br_false(self, arg: int) -> list[State]:
        condition = self.stack.pop()
        if isinstance(condition, bool):
            if condition is False:
                self.path.append("Jump")
                self.ip = arg
            else:
                self.path.append("No Jump")
                self.ip += 1
            return [self]
        else:
            taken = self
            not_taken = self.clone()
            taken.solver.constraints.append(z3.Not(condition))
            not_taken.solver.constraints.append(condition)

            taken.ip = arg
            not_taken.ip += 1
            taken.path.append("Jump")
            not_taken.path.append("No Jump")

            return [s for s in [taken, not_taken] if s.is_satisfiable()]

    def _st_loc(self, arg: int) -> list[State]:
        self.locals[arg] = self.stack.pop()
        self.ip += 1
        return [self]

    def _copy_loc(self, arg: int) -> list[State]:
        return self.__get_loc(arg)

    def _move_loc(self, arg: int) -> list[State]:
        return self.__get_loc(arg)

    def _mut_borrow_loc(self, arg: int) -> list[State]:
        return self.__get_loc(arg)

    def _imm_borrow_loc(self, arg: int) -> list[State]:
        return self.__get_loc(arg)

    def __get_loc(self, arg: int) -> int:
        self.stack.append(self.locals[arg])
        self.ip += 1
        return [self]

    def _ld_const(self, arg: int) -> list[State]:
        self.stack.append(CONST_ARR)
        self.ip += 1
        return [self]

    def _read_ref(self, arg: int) -> list[State]:
        self.ip += 1
        return [self]

    def _pop(self, arg: int) -> list[State]:
        self.stack.pop()
        self.ip += 1
        return [self]

    def _abort(self, arg: int) -> list[State]:
        # print(self.path)
        self.pc = 0x1337deadbeef
        return [self]

    def _ret(self, arg: int) -> list[State]:
        # print(self.path)
        raise NotImplementedError()

    def _vec_imm_borrow(self, arg: int) -> list[State]:
        idx = self.stack.pop()
        vec = self.stack.pop()
        assert isinstance(vec, list)
        assert isinstance(idx, int)
        self.stack.append(vec[idx])
        self.ip += 1
        return [self]

    def _vec_len(self, arg: int) -> list[State]:
        self.stack.append(len(self.stack.pop()))
        self.ip += 1
        return [self]

    def _vec_pop_back(self, arg: int) -> list[State]:
        vec = self.stack.pop()
        assert isinstance(vec, list)
        self.stack.append(vec.pop())
        self.ip += 1
        return [self]

    def _vec_push_back(self, arg: int) -> list[State]:
        # This can be why it's wrong
        elem = self.stack.pop()
        vec = self.stack.pop()

        assert isinstance(vec, list)
        vec.append(elem)
        self.ip += 1
        return [self]

    def _vec_pack(self, arg: int) -> list[State]:
        self.stack.append([self.stack.pop() for _ in range(arg)])
        self.ip += 1
        return [self]

    def _cast_u8(self, arg: int) -> list[State]:
        return self.__cast(8)

    def _cast_u64(self, arg: int) -> list[State]:
        return self.__cast(64)

    def _cast_u128(self, arg: int) -> list[State]:
        return self.__cast(128)

    def __cast(self, size: int):
        self.stack.append(self.stack.pop() & ((1 << size) - 1))
        self.ip += 1
        return [self]

    def _ld_u8(self, arg: int) -> list[State]:
        return self.__ld(arg)

    def _ld_u64(self, arg: int) -> list[State]:
        return self.__ld(arg)

    def _ld_u128(self, arg: int) -> list[State]:
        return self.__ld(arg)

    def __ld(self, arg: int) -> list[State]:
        self.stack.append(arg)
        self.ip += 1
        return [self]