diff --git a/dissect/cstruct/__init__.py b/dissect/cstruct/__init__.py index 9e01ee4..d6d78c1 100644 --- a/dissect/cstruct/__init__.py +++ b/dissect/cstruct/__init__.py @@ -30,12 +30,18 @@ ) from dissect.cstruct.util import ( dumpstruct, + f16, + f32, + f64, hexdump, p8, p16, p32, p64, pack, + pf16, + pf32, + pf64, swap, swap16, swap32, @@ -75,12 +81,18 @@ "ctypes", "ctypes_type", "dumpstruct", + "f16", + "f32", + "f64", "hexdump", "p8", "p16", "p32", "p64", "pack", + "pf16", + "pf32", + "pf64", "swap", "swap16", "swap32", diff --git a/dissect/cstruct/compiler.py b/dissect/cstruct/compiler.py index a096cad..aaf2e7b 100644 --- a/dissect/cstruct/compiler.py +++ b/dissect/cstruct/compiler.py @@ -11,6 +11,7 @@ from dissect.cstruct.bitbuffer import BitBuffer from dissect.cstruct.types import ( Array, + BaseArray, Char, CharArray, Flag, @@ -24,18 +25,15 @@ Wchar, WcharArray, ) -from dissect.cstruct.types.base import BaseArray from dissect.cstruct.types.enum import EnumMetaType -from dissect.cstruct.types.packed import _struct +from dissect.cstruct.util import _struct if TYPE_CHECKING: from collections.abc import Iterator from types import MethodType from dissect.cstruct.cstruct import cstruct - from dissect.cstruct.types import ( - BaseType, - ) + from dissect.cstruct.types import BaseType from dissect.cstruct.types.structure import Field SUPPORTED_TYPES = ( diff --git a/dissect/cstruct/types/base.py b/dissect/cstruct/types/base.py index 1c306f4..0127250 100644 --- a/dissect/cstruct/types/base.py +++ b/dissect/cstruct/types/base.py @@ -6,6 +6,7 @@ from dissect.cstruct.exception import ArraySizeError from dissect.cstruct.expression import Expression +from dissect.cstruct.util import normalize_endianness if TYPE_CHECKING: from collections.abc import Callable @@ -328,27 +329,5 @@ def _is_eof(stream: BinaryIO) -> bool: return False -ENDIANNESS_MAP: dict[AllowedEndianness, Endianness] = { - "<": "<", - ">": ">", - "!": "!", - "@": "@", - "=": "=", - "network": "!", - "little": "<", - "big": ">", -} - - -def normalize_endianness(endian: AllowedEndianness) -> Endianness: - """Normalize an endianness string to one of the standard format characters.""" - try: - return ENDIANNESS_MAP[endian] - except KeyError: - raise ValueError( - f"Invalid endianness: {endian!r}, expected one of {', '.join(ENDIANNESS_MAP.keys())}" - ) from None - - # As mentioned in the BaseType class, we correctly set the type here MetaType.ArrayType = Array diff --git a/dissect/cstruct/types/packed.py b/dissect/cstruct/types/packed.py index b095e28..6479ca4 100644 --- a/dissect/cstruct/types/packed.py +++ b/dissect/cstruct/types/packed.py @@ -1,10 +1,9 @@ from __future__ import annotations -from functools import lru_cache -from struct import Struct from typing import TYPE_CHECKING, Any, BinaryIO, Generic, TypeVar from dissect.cstruct.types.base import EOF, BaseType +from dissect.cstruct.util import _struct if TYPE_CHECKING: from typing_extensions import Self @@ -12,11 +11,6 @@ from dissect.cstruct.cstruct import Endianness -@lru_cache(1024) -def _struct(endian: str, packchar: str) -> Struct: - return Struct(f"{endian}{packchar}") - - T = TypeVar("T", int, float) diff --git a/dissect/cstruct/util.py b/dissect/cstruct/util.py index 00b5ad3..8e25f8f 100644 --- a/dissect/cstruct/util.py +++ b/dissect/cstruct/util.py @@ -3,15 +3,17 @@ import os import pprint import string +import struct import sys from enum import Enum +from functools import lru_cache from typing import TYPE_CHECKING if TYPE_CHECKING: from collections.abc import Iterator from typing import Literal - from dissect.cstruct.cstruct import AllowedEndianness + from dissect.cstruct.cstruct import AllowedEndianness, Endianness from dissect.cstruct.types.base import BaseType from dissect.cstruct.types.structure import Structure @@ -64,6 +66,28 @@ "big": "big", } +ENDIANNESS_MAP: dict[AllowedEndianness, Endianness] = { + "<": "<", + ">": ">", + "!": "!", + "@": "@", + "=": "=", + "network": "!", + "little": "<", + "big": ">", +} + + +def normalize_endianness(endian: AllowedEndianness) -> Endianness: + """Normalize an endianness string to one of the standard format characters.""" + try: + return ENDIANNESS_MAP[endian] + except KeyError: + raise ValueError( + f"Invalid endianness: {endian!r}, expected one of {', '.join(ENDIANNESS_MAP.keys())}" + ) from None + + Palette = list[tuple[int, str]] @@ -327,6 +351,11 @@ def dumpstruct( return _dumpstruct(obj, obj.dumps(), offset, color, output, autoskip) +@lru_cache(1024) +def _struct(endian: str, packchar: str) -> struct.Struct: + return struct.Struct(f"{endian}{packchar}") + + def pack(value: int, size: int | None = None, endian: AllowedEndianness = "little") -> bytes: """Pack an integer value to a given bit size, endianness. @@ -400,6 +429,36 @@ def p64(value: int, endian: AllowedEndianness = "little") -> bytes: return pack(value, 64, endian) +def pf16(value: float, endian: AllowedEndianness = "little") -> bytes: + """Pack a 16 bit float. + + Arguments: + value: Value to pack. + endian: Endianness to use (little, big, network, <, >, !, @ or =). + """ + return _struct(normalize_endianness(endian), "e").pack(value) + + +def pf32(value: float, endian: AllowedEndianness = "little") -> bytes: + """Pack a 32 bit float. + + Arguments: + value: Value to pack. + endian: Endianness to use (little, big, network, <, >, !, @ or =). + """ + return _struct(normalize_endianness(endian), "f").pack(value) + + +def pf64(value: float, endian: AllowedEndianness = "little") -> bytes: + """Pack a 64 bit float. + + Arguments: + value: Value to pack. + endian: Endianness to use (little, big, network, <, >, !, @ or =). + """ + return _struct(normalize_endianness(endian), "d").pack(value) + + def u8(value: bytes, endian: AllowedEndianness = "little", sign: bool = False) -> int: """Unpack an 8 bit integer. @@ -444,6 +503,36 @@ def u64(value: bytes, endian: AllowedEndianness = "little", sign: bool = False) return unpack(value, 64, endian, sign) +def f16(value: bytes, endian: AllowedEndianness = "little") -> float: + """Unpack a 16 bit float. + + Arguments: + value: Value to unpack. + endian: Endianness to use (little, big, network, <, >, !, @ or =). + """ + return _struct(normalize_endianness(endian), "e").unpack(value)[0] + + +def f32(value: bytes, endian: AllowedEndianness = "little") -> float: + """Unpack a 32 bit float. + + Arguments: + value: Value to unpack. + endian: Endianness to use (little, big, network, <, >, !, @ or =). + """ + return _struct(normalize_endianness(endian), "f").unpack(value)[0] + + +def f64(value: bytes, endian: AllowedEndianness = "little") -> float: + """Unpack a 64 bit float. + + Arguments: + value: Value to unpack. + endian: Endianness to use (little, big, network, <, >, !, @ or =). + """ + return _struct(normalize_endianness(endian), "d").unpack(value)[0] + + def swap(value: int, size: int) -> int: """Swap the endianness of an integer with a given bit size. diff --git a/tests/test_utils.py b/tests/test_utils.py index 32ac64b..60f7d67 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -259,10 +259,16 @@ def test_pack_unpack() -> None: assert util.p16(1, endian) == b"\x01\x00" assert util.p32(1, endian) == b"\x01\x00\x00\x00" assert util.p64(1, endian) == b"\x01\x00\x00\x00\x00\x00\x00\x00" + assert util.pf16(1.0, endian) == b"\x00\x3c" + assert util.pf32(1.0, endian) == b"\x00\x00\x80\x3f" + assert util.pf64(1.0, endian) == b"\x00\x00\x00\x00\x00\x00\xf0\x3f" assert util.u8(b"\x01", endian, sign) == 1 assert util.u16(b"\x01\x00", endian, sign) == 1 assert util.u32(b"\x01\x00\x00\x00", endian, sign) == 1 assert util.u64(b"\x01\x00\x00\x00\x00\x00\x00\x00", endian, sign) == 1 + assert util.f16(b"\x00\x3c", endian) == 1.0 + assert util.f32(b"\x00\x00\x80\x3f", endian) == 1.0 + assert util.f64(b"\x00\x00\x00\x00\x00\x00\xf0\x3f", endian) == 1.0 endian = "big" sign = False @@ -270,10 +276,16 @@ def test_pack_unpack() -> None: assert util.p16(1, endian) == b"\x00\x01" assert util.p32(1, endian) == b"\x00\x00\x00\x01" assert util.p64(1, endian) == b"\x00\x00\x00\x00\x00\x00\x00\x01" + assert util.pf16(1.0, endian) == b"\x3c\x00" + assert util.pf32(1.0, endian) == b"\x3f\x80\x00\x00" + assert util.pf64(1.0, endian) == b"\x3f\xf0\x00\x00\x00\x00\x00\x00" assert util.u8(b"\x01", endian, sign) == 1 assert util.u16(b"\x00\x01", endian, sign) == 1 assert util.u32(b"\x00\x00\x00\x01", endian, sign) == 1 assert util.u64(b"\x00\x00\x00\x00\x00\x00\x00\x01", endian, sign) == 1 + assert util.f16(b"\x3c\x00", endian) == 1.0 + assert util.f32(b"\x3f\x80\x00\x00", endian) == 1.0 + assert util.f64(b"\x3f\xf0\x00\x00\x00\x00\x00\x00", endian) == 1.0 endian = "network" sign = False @@ -281,10 +293,16 @@ def test_pack_unpack() -> None: assert util.p16(1, endian) == b"\x00\x01" assert util.p32(1, endian) == b"\x00\x00\x00\x01" assert util.p64(1, endian) == b"\x00\x00\x00\x00\x00\x00\x00\x01" + assert util.pf16(1.0, endian) == b"\x3c\x00" + assert util.pf32(1.0, endian) == b"\x3f\x80\x00\x00" + assert util.pf64(1.0, endian) == b"\x3f\xf0\x00\x00\x00\x00\x00\x00" assert util.u8(b"\x01", endian, sign) == 1 assert util.u16(b"\x00\x01", endian, sign) == 1 assert util.u32(b"\x00\x00\x00\x01", endian, sign) == 1 assert util.u64(b"\x00\x00\x00\x00\x00\x00\x00\x01", endian, sign) == 1 + assert util.f16(b"\x3c\x00", endian) == 1.0 + assert util.f32(b"\x3f\x80\x00\x00", endian) == 1.0 + assert util.f64(b"\x3f\xf0\x00\x00\x00\x00\x00\x00", endian) == 1.0 endian = "little" sign = True @@ -292,10 +310,16 @@ def test_pack_unpack() -> None: assert util.p16(-120, endian) == b"\x88\xff" assert util.p32(-120, endian) == b"\x88\xff\xff\xff" assert util.p64(-120, endian) == b"\x88\xff\xff\xff\xff\xff\xff\xff" + assert util.pf16(-120.0, endian) == b"\x80\xd7" + assert util.pf32(-120.0, endian) == b"\x00\x00\xf0\xc2" + assert util.pf64(-120.0, endian) == b"\x00\x00\x00\x00\x00\x00\x5e\xc0" assert util.u8(b"\x88", endian, sign) == -120 assert util.u16(b"\x88\xff", endian, sign) == -120 assert util.u32(b"\x88\xff\xff\xff", endian, sign) == -120 assert util.u64(b"\x88\xff\xff\xff\xff\xff\xff\xff", endian, sign) == -120 + assert util.f16(b"\x80\xd7", endian) == -120.0 + assert util.f32(b"\x00\x00\xf0\xc2", endian) == -120.0 + assert util.f64(b"\x00\x00\x00\x00\x00\x00\x5e\xc0", endian) == -120.0 endian = "big" sign = True @@ -303,10 +327,16 @@ def test_pack_unpack() -> None: assert util.p16(-120, endian) == b"\xff\x88" assert util.p32(-120, endian) == b"\xff\xff\xff\x88" assert util.p64(-120, endian) == b"\xff\xff\xff\xff\xff\xff\xff\x88" + assert util.pf16(-120.0, endian) == b"\xd7\x80" + assert util.pf32(-120.0, endian) == b"\xc2\xf0\x00\x00" + assert util.pf64(-120.0, endian) == b"\xc0\x5e\x00\x00\x00\x00\x00\x00" assert util.u8(b"\x88", endian, sign) == -120 assert util.u16(b"\xff\x88", endian, sign) == -120 assert util.u32(b"\xff\xff\xff\x88", endian, sign) == -120 assert util.u64(b"\xff\xff\xff\xff\xff\xff\xff\x88", endian, sign) == -120 + assert util.f16(b"\xd7\x80", endian) == -120.0 + assert util.f32(b"\xc2\xf0\x00\x00", endian) == -120.0 + assert util.f64(b"\xc0\x5e\x00\x00\x00\x00\x00\x00", endian) == -120.0 assert util.pack(1, 24) == b"\x01\x00\x00" assert util.unpack(b"\x01\x00\x00", 24) == 1