Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/dissect/cstruct/utils.py: 29%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

153 statements  

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)