145 lines
5.0 KiB
Python
145 lines
5.0 KiB
Python
#-------------------------------------------------------------------------------
|
|
# elftools: common/construct_utils.py
|
|
#
|
|
# Some complementary construct utilities
|
|
#
|
|
# Eli Bendersky (eliben@gmail.com)
|
|
# This code is in the public domain
|
|
#-------------------------------------------------------------------------------
|
|
from __future__ import annotations
|
|
|
|
from struct import Struct
|
|
from typing import IO, TYPE_CHECKING, Any, NoReturn
|
|
|
|
from ..construct import (
|
|
Subconstruct, ConstructError, ArrayError, SizeofError, Construct, StaticField, FieldError
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable, Iterable
|
|
|
|
from ..construct import Container
|
|
|
|
|
|
class RepeatUntilExcluding(Subconstruct):
|
|
""" A version of construct's RepeatUntil that doesn't include the last
|
|
element (which casued the repeat to exit) in the return value.
|
|
|
|
Only parsing is currently implemented.
|
|
|
|
P.S. removed some code duplication
|
|
"""
|
|
__slots__ = ("predicate",)
|
|
def __init__(self, predicate: Callable[[Any, Container], bool], subcon: Construct) -> None:
|
|
Subconstruct.__init__(self, subcon)
|
|
self.predicate = predicate
|
|
self._clear_flag(self.FLAG_COPY_CONTEXT)
|
|
self._set_flag(self.FLAG_DYNAMIC)
|
|
def _parse(self, stream: IO[bytes], context: Container) -> list[Any]:
|
|
obj = []
|
|
try:
|
|
context_for_subcon = context
|
|
if self.subcon.conflags & self.FLAG_COPY_CONTEXT:
|
|
context_for_subcon = context.__copy__()
|
|
|
|
while True:
|
|
subobj = self.subcon._parse(stream, context_for_subcon)
|
|
if self.predicate(subobj, context):
|
|
break
|
|
obj.append(subobj)
|
|
except ConstructError as ex:
|
|
raise ArrayError("missing terminator", ex)
|
|
return obj
|
|
def _build(self, obj: Iterable[Any], stream: IO[bytes], context: Container) -> NoReturn:
|
|
raise NotImplementedError('no building')
|
|
def _sizeof(self, context: Container) -> int:
|
|
raise SizeofError("can't calculate size")
|
|
|
|
|
|
class _NamedConstruct(Construct):
|
|
if TYPE_CHECKING:
|
|
name: str # instead of `str|None` from Construct to save us from `is None` checks everywhere
|
|
|
|
|
|
class ULEB128(_NamedConstruct):
|
|
"""A construct based parser for ULEB128 encoding.
|
|
"""
|
|
def _parse(self, stream: IO[bytes], context: Container) -> int:
|
|
value = 0
|
|
shift = 0
|
|
while True:
|
|
data = stream.read(1)
|
|
if len(data) != 1:
|
|
raise FieldError("unexpected end of stream while parsing a ULEB128 encoded value")
|
|
b = data[0]
|
|
value |= (b & 0x7F) << shift
|
|
shift += 7
|
|
if b & 0x80 == 0:
|
|
return value
|
|
|
|
|
|
class SLEB128(_NamedConstruct):
|
|
"""A construct based parser for SLEB128 encoding.
|
|
"""
|
|
def _parse(self, stream: IO[bytes], context: Container) -> int:
|
|
value = 0
|
|
shift = 0
|
|
while True:
|
|
data = stream.read(1)
|
|
if len(data) != 1:
|
|
raise FieldError("unexpected end of stream while parsing a SLEB128 encoded value")
|
|
b = data[0]
|
|
value |= (b & 0x7F) << shift
|
|
shift += 7
|
|
if b & 0x80 == 0:
|
|
return value | (~0 << shift) if b & 0x40 else value
|
|
|
|
|
|
class StreamOffset(_NamedConstruct):
|
|
"""
|
|
Captures the current stream offset
|
|
|
|
Parameters:
|
|
* name - the name of the value
|
|
|
|
Example:
|
|
StreamOffset("item_offset")
|
|
"""
|
|
__slots__ = ()
|
|
def __init__(self, name: str) -> None:
|
|
Construct.__init__(self, name)
|
|
self._set_flag(self.FLAG_DYNAMIC)
|
|
def _parse(self, stream: IO[bytes], context: Container) -> int:
|
|
return stream.tell()
|
|
def _build(self, obj: None, stream: IO[bytes], context: Container) -> None:
|
|
context[self.name] = stream.tell()
|
|
def _sizeof(self, context: Container) -> int:
|
|
return 0
|
|
|
|
_UBInt24_packer = Struct(">BH")
|
|
_ULInt24_packer = Struct("<HB")
|
|
|
|
class UBInt24(StaticField):
|
|
"""unsigned, big endian 24-bit integer"""
|
|
def __init__(self, name: str) -> None:
|
|
StaticField.__init__(self, name, 3)
|
|
|
|
def _parse(self, stream: IO[bytes], context: Container) -> int:
|
|
(h, l) = _UBInt24_packer.unpack(StaticField._parse(self, stream, context))
|
|
return l | (h << 16)
|
|
|
|
def _build(self, obj: int, stream: IO[bytes], context: Container) -> None:
|
|
StaticField._build(self, _UBInt24_packer.pack(obj >> 16, obj & 0xFFFF), stream, context)
|
|
|
|
class ULInt24(StaticField):
|
|
"""unsigned, little endian 24-bit integer"""
|
|
def __init__(self, name: str) -> None:
|
|
StaticField.__init__(self, name, 3)
|
|
|
|
def _parse(self, stream: IO[bytes], context: Container) -> int:
|
|
(l, h) = _ULInt24_packer.unpack(StaticField._parse(self, stream, context))
|
|
return l | (h << 16)
|
|
|
|
def _build(self, obj: int, stream: IO[bytes], context: Container) -> None:
|
|
StaticField._build(self, _ULInt24_packer.pack(obj & 0xFFFF, obj >> 16), stream, context)
|