1from __future__ import annotations
2
3import pprint
4import string
5import sys
6from enum import Enum
7from typing import TYPE_CHECKING
8
9from dissect.cstruct.types.pointer import Pointer
10from dissect.cstruct.types.structure import Structure
11
12if TYPE_CHECKING:
13 from collections.abc import Iterator
14 from typing import Literal
15
16COLOR_RED = "\033[1;31m"
17COLOR_GREEN = "\033[1;32m"
18COLOR_YELLOW = "\033[1;33m"
19COLOR_BLUE = "\033[1;34m"
20COLOR_PURPLE = "\033[1;35m"
21COLOR_CYAN = "\033[1;36m"
22COLOR_WHITE = "\033[1;37m"
23COLOR_NORMAL = "\033[1;0m"
24
25COLOR_BG_RED = "\033[1;41m\033[1;37m"
26COLOR_BG_GREEN = "\033[1;42m\033[1;37m"
27COLOR_BG_YELLOW = "\033[1;43m\033[1;37m"
28COLOR_BG_BLUE = "\033[1;44m\033[1;37m"
29COLOR_BG_PURPLE = "\033[1;45m\033[1;37m"
30COLOR_BG_CYAN = "\033[1;46m\033[1;37m"
31COLOR_BG_WHITE = "\033[1;47m\033[1;30m"
32
33PRINTABLE = string.digits + string.ascii_letters + string.punctuation + " "
34
35ENDIANNESS_MAP: dict[str, Literal["big", "little"]] = {
36 "@": sys.byteorder,
37 "=": sys.byteorder,
38 "<": "little",
39 ">": "big",
40 "!": "big",
41 "network": "big",
42}
43
44Palette = list[tuple[int, str]]
45
46
47def _hexdump(data: bytes, palette: Palette | None = None, offset: int = 0, prefix: str = "") -> Iterator[str]:
48 """Hexdump some data.
49
50 Args:
51 data: Bytes to hexdump.
52 offset: Byte offset of the hexdump.
53 prefix: Optional prefix.
54 palette: Colorize the hexdump using this color pattern.
55 """
56 if palette:
57 palette = palette[::-1]
58
59 remaining = 0
60 active = None
61
62 for i in range(0, len(data), 16):
63 values = ""
64 chars = []
65
66 for j in range(16):
67 if not active and palette:
68 remaining, active = palette.pop()
69 while remaining == 0:
70 if len(palette) == 0:
71 # Last palette tuple is empty: print remaining whitespaces
72 active = ""
73 break
74 else:
75 remaining, active = palette.pop()
76 values += active
77 elif active and j == 0:
78 values += active
79
80 if i + j >= len(data):
81 values += " "
82 else:
83 char = data[i + j]
84 char = chr(char)
85
86 print_char = char if char in PRINTABLE else "."
87
88 if active:
89 values += f"{ord(char):02x}"
90 chars.append(active + print_char + COLOR_NORMAL)
91 else:
92 values += f"{ord(char):02x}"
93 chars.append(print_char)
94
95 remaining -= 1
96 if remaining == 0:
97 active = None
98
99 if palette is not None:
100 values += COLOR_NORMAL
101
102 if j == 15 and palette is not None:
103 values += COLOR_NORMAL
104
105 values += " "
106 if j == 7:
107 values += " "
108
109 chars = "".join(chars)
110 yield f"{prefix}{offset + i:08x} {values:48s} {chars}"
111
112
113def hexdump(
114 data: bytes, palette: Palette | None = None, offset: int = 0, prefix: str = "", output: str = "print"
115) -> Iterator[str] | str | None:
116 """Hexdump some data.
117
118 Args:
119 data: Bytes to hexdump.
120 palette: Colorize the hexdump using this color pattern.
121 offset: Byte offset of the hexdump.
122 prefix: Optional prefix.
123 output: Output format, can be 'print', 'generator' or 'string'.
124 """
125 generator = _hexdump(data, palette, offset, prefix)
126 if output == "print":
127 print("\n".join(generator))
128 return None
129 if output == "generator":
130 return generator
131 if output == "string":
132 return "\n".join(list(generator))
133 raise ValueError(f"Invalid output argument: {output!r} (should be 'print', 'generator' or 'string').")
134
135
136def _dumpstruct(
137 structure: Structure,
138 data: bytes,
139 offset: int,
140 color: bool,
141 output: str,
142) -> str | None:
143 palette = []
144 colors = [
145 (COLOR_RED, COLOR_BG_RED),
146 (COLOR_GREEN, COLOR_BG_GREEN),
147 (COLOR_YELLOW, COLOR_BG_YELLOW),
148 (COLOR_BLUE, COLOR_BG_BLUE),
149 (COLOR_PURPLE, COLOR_BG_PURPLE),
150 (COLOR_CYAN, COLOR_BG_CYAN),
151 (COLOR_WHITE, COLOR_BG_WHITE),
152 ]
153 ci = 0
154 out = [f"struct {structure.__class__.__name__}:"]
155 foreground, background = None, None
156 for field in structure.__class__.__fields__:
157 if getattr(field.type, "anonymous", False):
158 continue
159
160 value = getattr(structure, field._name)
161 if isinstance(value, (str, Pointer, Enum)):
162 value = repr(value)
163 elif isinstance(value, int):
164 value = hex(value)
165 elif isinstance(value, list):
166 value = pprint.pformat(value)
167 if "\n" in value:
168 value = value.replace("\n", f"\n{' ' * (len(field._name) + 4)}")
169
170 if color:
171 foreground, background = colors[ci % len(colors)]
172 size = structure.__sizes__[field._name]
173 palette.append((size, background))
174 ci += 1
175 out.append(f"- {foreground}{field._name}{COLOR_NORMAL}: {value}")
176 else:
177 out.append(f"- {field._name}: {value}")
178
179 out = "\n".join(out)
180
181 if output == "print":
182 print()
183 hexdump(data, palette, offset=offset)
184 print()
185 print(out)
186 elif output == "string":
187 return f"\n{hexdump(data, palette, offset=offset, output='string')}\n\n{out}"
188 return None
189
190
191def dumpstruct(
192 obj: Structure | type[Structure],
193 data: bytes | None = None,
194 offset: int = 0,
195 color: bool = True,
196 output: str = "print",
197) -> str | None:
198 """Dump a structure or parsed structure instance.
199
200 Prints a colorized hexdump and parsed structure output.
201
202 Args:
203 obj: Structure to dump.
204 data: Bytes to parse the Structure on, if obj is not a parsed Structure already.
205 offset: Byte offset of the hexdump.
206 output: Output format, can be 'print' or 'string'.
207 """
208 if output not in ("print", "string"):
209 raise ValueError(f"Invalid output argument: {output!r} (should be 'print' or 'string').")
210
211 if isinstance(obj, Structure):
212 return _dumpstruct(obj, obj.dumps(), offset, color, output)
213 if issubclass(obj, Structure) and data is not None:
214 return _dumpstruct(obj(data), data, offset, color, output)
215 raise ValueError("Invalid arguments")
216
217
218def pack(value: int, size: int | None = None, endian: str = "little") -> bytes:
219 """Pack an integer value to a given bit size, endianness.
220
221 Arguments:
222 value: Value to pack.
223 size: Integer size in bits.
224 endian: Endianness to use (little, big, network, <, > or !)
225 """
226 size = ((size or value.bit_length()) + 7) // 8
227 return value.to_bytes(size, ENDIANNESS_MAP.get(endian, endian), signed=value < 0)
228
229
230def unpack(value: bytes, size: int | None = None, endian: str = "little", sign: bool = False) -> int:
231 """Unpack an integer value from a given bit size, endianness and sign.
232
233 Arguments:
234 value: Value to unpack.
235 size: Integer size in bits.
236 endian: Endianness to use (little, big, network, <, > or !)
237 sign: Signedness of the integer.
238 """
239 if size and len(value) != size // 8:
240 raise ValueError(f"Invalid byte value, expected {size // 8} bytes, got {len(value)} bytes")
241 return int.from_bytes(value, ENDIANNESS_MAP.get(endian, endian), signed=sign)
242
243
244def p8(value: int, endian: str = "little") -> bytes:
245 """Pack an 8 bit integer.
246
247 Arguments:
248 value: Value to pack.
249 endian: Endianness to use (little, big, network, <, > or !)
250 """
251 return pack(value, 8, endian)
252
253
254def p16(value: int, endian: str = "little") -> bytes:
255 """Pack a 16 bit integer.
256
257 Arguments:
258 value: Value to pack.
259 endian: Endianness to use (little, big, network, <, > or !)
260 """
261 return pack(value, 16, endian)
262
263
264def p32(value: int, endian: str = "little") -> bytes:
265 """Pack a 32 bit integer.
266
267 Arguments:
268 value: Value to pack.
269 endian: Endianness to use (little, big, network, <, > or !)
270 """
271 return pack(value, 32, endian)
272
273
274def p64(value: int, endian: str = "little") -> bytes:
275 """Pack a 64 bit integer.
276
277 Arguments:
278 value: Value to pack.
279 endian: Endianness to use (little, big, network, <, > or !)
280 """
281 return pack(value, 64, endian)
282
283
284def u8(value: bytes, endian: str = "little", sign: bool = False) -> int:
285 """Unpack an 8 bit integer.
286
287 Arguments:
288 value: Value to unpack.
289 endian: Endianness to use (little, big, network, <, > or !)
290 sign: Signedness of the integer.
291 """
292 return unpack(value, 8, endian, sign)
293
294
295def u16(value: bytes, endian: str = "little", sign: bool = False) -> int:
296 """Unpack a 16 bit integer.
297
298 Arguments:
299 value: Value to unpack.
300 endian: Endianness to use (little, big, network, <, > or !)
301 sign: Signedness of the integer.
302 """
303 return unpack(value, 16, endian, sign)
304
305
306def u32(value: bytes, endian: str = "little", sign: bool = False) -> int:
307 """Unpack a 32 bit integer.
308
309 Arguments:
310 value: Value to unpack.
311 endian: Endianness to use (little, big, network, <, > or !)
312 sign: Signedness of the integer.
313 """
314 return unpack(value, 32, endian, sign)
315
316
317def u64(value: bytes, endian: str = "little", sign: bool = False) -> int:
318 """Unpack a 64 bit integer.
319
320 Arguments:
321 value: Value to unpack.
322 endian: Endianness to use (little, big, network, <, > or !)
323 sign: Signedness of the integer.
324 """
325 return unpack(value, 64, endian, sign)
326
327
328def swap(value: int, size: int) -> int:
329 """Swap the endianness of an integer with a given bit size.
330
331 Arguments:
332 value: Integer to swap.
333 size: Integer size in bits.
334 """
335 return unpack(pack(value, size, ">"), size, "<")
336
337
338def swap16(value: int) -> int:
339 """Swap the endianness of a 16 bit integer.
340
341 Arguments:
342 value: Integer to swap.
343 """
344 return swap(value, 16)
345
346
347def swap32(value: int) -> int:
348 """Swap the endianness of a 32 bit integer.
349
350 Arguments:
351 value: Integer to swap.
352 """
353 return swap(value, 32)
354
355
356def swap64(value: int) -> int:
357 """Swap the endianness of a 64 bit integer.
358
359 Arguments:
360 value: Integer to swap.
361 """
362 return swap(value, 64)