from __future__ import annotations
from enum import Enum, auto
import claripy
import ailment
from archinfo import Arch, RegisterOffset
from angr.calling_conventions import SimFunctionArgument, SimRegArg, SimStackArg
from angr.engines.light import SpOffset
from .heap_address import HeapAddress
[文档]
class AtomKind(Enum):
"""
An enum indicating the class of an atom
"""
REGISTER = auto()
MEMORY = auto()
TMP = auto()
GUARD = auto()
CONSTANT = auto()
[文档]
class Atom:
"""
This class represents a data storage location manipulated by IR instructions.
It could either be a Tmp (temporary variable), a Register, a MemoryLocation.
"""
__slots__ = ("_hash", "size")
[文档]
def __init__(self, size):
"""
:param size: The size of the atom in bytes
"""
self.size = size
self._hash = None
def __repr__(self):
raise NotImplementedError
@property
def bits(self) -> int:
return self.size * 8
@property
def _size(self):
return self.size
@_size.setter
def _size(self, v):
self.size = v
[文档]
@staticmethod
def from_ail_expr(expr: ailment.Expr.Expression, arch: Arch, full_reg: bool = False) -> Register:
if isinstance(expr, ailment.Expr.Register):
if full_reg:
reg_name = arch.translate_register_name(expr.reg_offset)
return Register(arch.registers[reg_name][0], arch.registers[reg_name][1], arch)
return Register(expr.reg_offset, expr.size, arch)
raise TypeError(f"Expression type {type(expr)} is not yet supported")
[文档]
@staticmethod
def from_argument(
argument: SimFunctionArgument, arch: Arch, full_reg=False, sp: int | None = None
) -> Register | MemoryLocation:
"""
Instantiate an `Atom` from a given argument.
:param argument: The argument to create a new atom from.
:param arch: The argument representing archinfo architecture for argument.
:param full_reg: Whether to return an atom indicating the entire register if the argument only specifies a
slice of the register.
:param sp: The current stack offset. Optional. Only used when argument is a SimStackArg.
"""
if isinstance(argument, SimRegArg):
if full_reg:
return Register(arch.registers[argument.reg_name][0], arch.registers[argument.reg_name][1], arch)
return Register(arch.registers[argument.reg_name][0] + argument.reg_offset, argument.size, arch)
if isinstance(argument, SimStackArg):
if sp is None:
raise ValueError("You must provide a stack pointer to translate a SimStackArg")
return MemoryLocation(
SpOffset(arch.bits, argument.stack_offset + sp), argument.size, endness=arch.memory_endness
)
raise TypeError(f"Argument type {type(argument)} is not yet supported.")
[文档]
@staticmethod
def reg(thing: str | RegisterOffset, size: int | None = None, arch: Arch | None = None) -> Register:
"""
Create a Register atom.
:param thing: The register offset (e.g., project.arch.registers["rax"][0]) or the register name (e.g., "rax").
:param size: Size of the register atom. Must be provided when creating the atom using a register offset.
:param arch: The architecture. Must be provided when creating the atom using a register name.
:return: The Register Atom object.
"""
if isinstance(thing, str):
if arch is None:
raise ValueError(
"Cannot create a Register Atom by register name without having an architecture "
"specified through arch!"
)
if thing not in arch.registers:
raise ValueError(f"Unknown register name {thing} for architecture {arch.name}")
reg_offset, size_ = arch.registers[thing]
if size is None:
size = size_
elif isinstance(thing, RegisterOffset):
reg_offset = thing
if size is None:
raise ValueError("You must provide a size when specifying the register offset")
else:
raise TypeError(
"Unsupported type of register. It must be a string (for register name) or an int (for "
"register offset)"
)
return Register(reg_offset, size, arch=arch)
register = reg
[文档]
@staticmethod
def mem(addr: SpOffset | HeapAddress | int, size: int, endness: str | None = None) -> MemoryLocation:
"""
Create a MemoryLocation atom,
:param addr: The memory location. Can be an SpOffset for stack variables, an int for global memory
variables, or a HeapAddress for items on the heap.
:param size: Size of the atom.
:param endness: Optional, either "Iend_LE" or "Iend_BE".
:return: The MemoryLocation Atom object.
"""
return MemoryLocation(addr, size, endness=endness)
memory = mem
def _identity(self):
raise NotImplementedError
def __hash__(self):
if self._hash is None:
self._hash = hash(self._identity())
return self._hash
def __eq__(self, other):
return type(self) is type(other) and self._identity() == other._identity()
[文档]
class GuardUse(Atom):
"""
Implements a guard use.
"""
__slots__ = ("target",)
[文档]
def __init__(self, target):
super().__init__(0)
self.target = target
def __repr__(self):
return f"<Guard {self.target:#x}>"
def _identity(self):
return (self.target,)
[文档]
class ConstantSrc(Atom):
"""
Represents a constant.
"""
__slots__ = ("value",)
[文档]
def __init__(self, value: int, size: int):
super().__init__(size)
self.value: int = value
def __repr__(self):
return f"<Const {self.value}>"
def _identity(self):
return (self.value, self.size)
[文档]
class Tmp(Atom):
"""
Represents a variable used by the IR to store intermediate values.
"""
__slots__ = ("tmp_idx",)
[文档]
def __init__(self, tmp_idx: int, size: int):
super().__init__(size)
self.tmp_idx = tmp_idx
def __repr__(self):
return f"<Tmp {self.tmp_idx}>"
def _identity(self):
return hash(("tmp", self.tmp_idx))
[文档]
class Register(Atom):
"""
Represents a given CPU register.
As an IR abstracts the CPU design to target different architectures, registers are represented as a separated memory
space.
Thus a register is defined by its offset from the base of this memory and its size.
:ivar int reg_offset: The offset from the base to define its place in the memory bloc.
:ivar int size: The size, in number of bytes.
"""
__slots__ = (
"arch",
"reg_offset",
)
[文档]
def __init__(self, reg_offset: RegisterOffset | int, size: int, arch: Arch | None = None):
super().__init__(size)
self.reg_offset = RegisterOffset(reg_offset)
self.arch = arch
def __repr__(self):
return f"<Reg {self.name}<{self.size}>>"
def _identity(self):
return (self.reg_offset, self.size)
@property
def name(self) -> str:
return (
str(self.reg_offset) if self.arch is None else self.arch.translate_register_name(self.reg_offset, self.size)
)
[文档]
class VirtualVariable(Atom):
"""
Represents a virtual variable.
"""
__slots__ = (
"category",
"oident",
"varid",
)
[文档]
def __init__(
self, varid: int, size: int, category: ailment.Expr.VirtualVariableCategory, oident: str | int | None = None
):
super().__init__(size)
self.varid = varid
self.category = category
self.oident = oident
def __repr__(self):
return f"<VVar {self.varid}<{self.size}>>"
def _identity(self):
return self.varid, self.size
@property
def was_reg(self) -> bool:
return self.category == ailment.Expr.VirtualVariableCategory.REGISTER
@property
def was_stack(self) -> bool:
return self.category == ailment.Expr.VirtualVariableCategory.STACK
@property
def was_parameter(self) -> bool:
return self.category == ailment.Expr.VirtualVariableCategory.PARAMETER
@property
def was_tmp(self) -> bool:
return self.category == ailment.Expr.VirtualVariableCategory.TMP
@property
def reg_offset(self) -> int | None:
if self.was_reg:
return self.oident
return None
@property
def stack_offset(self) -> int | None:
if self.was_stack:
return self.oident
return None
@property
def tmp_idx(self) -> int | None:
return self.oident if self.was_tmp else None
[文档]
class MemoryLocation(Atom):
"""
Represents a memory slice.
It is characterized by its address and its size.
"""
__slots__ = (
"addr",
"endness",
)
[文档]
def __init__(self, addr: SpOffset | HeapAddress | int, size: int, endness: str | None = None):
"""
:param int addr: The address of the beginning memory location slice.
:param int size: The size of the represented memory location, in bytes.
"""
super().__init__(size)
self.addr: SpOffset | int | claripy.ast.BV = addr
self.endness = endness
def __repr__(self):
address_format = hex(self.addr) if type(self.addr) is int else self.addr
stack_format = " (stack)" if self.is_on_stack else ""
size = f"{self.size}" if isinstance(self.size, int) else self.size
return f"<Mem {address_format}<{size}>{stack_format}>"
@property
def is_on_stack(self) -> bool:
"""
True if this memory location is located on the stack.
"""
return isinstance(self.addr, SpOffset)
@property
def symbolic(self) -> bool:
if isinstance(self.addr, int):
return False
if isinstance(self.addr, SpOffset):
return type(self.addr.offset) is not int
return True
def __eq__(self, other):
# pylint:disable=isinstance-second-argument-not-valid-type
return (
type(other) is MemoryLocation
and (
self.addr is other.addr
if (isinstance(self.addr, claripy.ast.BV) or isinstance(other.addr, claripy.ast.BV))
else self.addr == other.addr
)
and self.size == other.size
and self.endness == other.endness
)
__hash__ = Atom.__hash__
def _identity(self):
return self.addr, self.size, self.endness
atom_kind_mapping = {
AtomKind.REGISTER: Register,
AtomKind.MEMORY: MemoryLocation,
AtomKind.TMP: Tmp,
AtomKind.GUARD: GuardUse,
AtomKind.CONSTANT: ConstantSrc,
}