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

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 if color: 

161 foreground, background = colors[ci % len(colors)] 

162 palette.append((structure._sizes[field._name], background)) 

163 ci += 1 

164 

165 value = getattr(structure, field._name) 

166 if isinstance(value, (str, Pointer, Enum)): 

167 value = repr(value) 

168 elif isinstance(value, int): 

169 value = hex(value) 

170 elif isinstance(value, list): 

171 value = pprint.pformat(value) 

172 if "\n" in value: 

173 value = value.replace("\n", f"\n{' ' * (len(field._name) + 4)}") 

174 

175 if color: 

176 out.append(f"- {foreground}{field._name}{COLOR_NORMAL}: {value}") 

177 else: 

178 out.append(f"- {field._name}: {value}") 

179 

180 out = "\n".join(out) 

181 

182 if output == "print": 

183 print() 

184 hexdump(data, palette, offset=offset) 

185 print() 

186 print(out) 

187 elif output == "string": 

188 return f"\n{hexdump(data, palette, offset=offset, output='string')}\n\n{out}" 

189 return None 

190 

191 

192def dumpstruct( 

193 obj: Structure | type[Structure], 

194 data: bytes | None = None, 

195 offset: int = 0, 

196 color: bool = True, 

197 output: str = "print", 

198) -> str | None: 

199 """Dump a structure or parsed structure instance. 

200 

201 Prints a colorized hexdump and parsed structure output. 

202 

203 Args: 

204 obj: Structure to dump. 

205 data: Bytes to parse the Structure on, if obj is not a parsed Structure already. 

206 offset: Byte offset of the hexdump. 

207 output: Output format, can be 'print' or 'string'. 

208 """ 

209 if output not in ("print", "string"): 

210 raise ValueError(f"Invalid output argument: {output!r} (should be 'print' or 'string').") 

211 

212 if isinstance(obj, Structure): 

213 return _dumpstruct(obj, obj.dumps(), offset, color, output) 

214 if issubclass(obj, Structure) and data is not None: 

215 return _dumpstruct(obj(data), data, offset, color, output) 

216 raise ValueError("Invalid arguments") 

217 

218 

219def pack(value: int, size: int | None = None, endian: str = "little") -> bytes: 

220 """Pack an integer value to a given bit size, endianness. 

221 

222 Arguments: 

223 value: Value to pack. 

224 size: Integer size in bits. 

225 endian: Endianness to use (little, big, network, <, > or !) 

226 """ 

227 size = ((size or value.bit_length()) + 7) // 8 

228 return value.to_bytes(size, ENDIANNESS_MAP.get(endian, endian), signed=value < 0) 

229 

230 

231def unpack(value: bytes, size: int | None = None, endian: str = "little", sign: bool = False) -> int: 

232 """Unpack an integer value from a given bit size, endianness and sign. 

233 

234 Arguments: 

235 value: Value to unpack. 

236 size: Integer size in bits. 

237 endian: Endianness to use (little, big, network, <, > or !) 

238 sign: Signedness of the integer. 

239 """ 

240 if size and len(value) != size // 8: 

241 raise ValueError(f"Invalid byte value, expected {size // 8} bytes, got {len(value)} bytes") 

242 return int.from_bytes(value, ENDIANNESS_MAP.get(endian, endian), signed=sign) 

243 

244 

245def p8(value: int, endian: str = "little") -> bytes: 

246 """Pack an 8 bit integer. 

247 

248 Arguments: 

249 value: Value to pack. 

250 endian: Endianness to use (little, big, network, <, > or !) 

251 """ 

252 return pack(value, 8, endian) 

253 

254 

255def p16(value: int, endian: str = "little") -> bytes: 

256 """Pack a 16 bit integer. 

257 

258 Arguments: 

259 value: Value to pack. 

260 endian: Endianness to use (little, big, network, <, > or !) 

261 """ 

262 return pack(value, 16, endian) 

263 

264 

265def p32(value: int, endian: str = "little") -> bytes: 

266 """Pack a 32 bit integer. 

267 

268 Arguments: 

269 value: Value to pack. 

270 endian: Endianness to use (little, big, network, <, > or !) 

271 """ 

272 return pack(value, 32, endian) 

273 

274 

275def p64(value: int, endian: str = "little") -> bytes: 

276 """Pack a 64 bit integer. 

277 

278 Arguments: 

279 value: Value to pack. 

280 endian: Endianness to use (little, big, network, <, > or !) 

281 """ 

282 return pack(value, 64, endian) 

283 

284 

285def u8(value: bytes, endian: str = "little", sign: bool = False) -> int: 

286 """Unpack an 8 bit integer. 

287 

288 Arguments: 

289 value: Value to unpack. 

290 endian: Endianness to use (little, big, network, <, > or !) 

291 sign: Signedness of the integer. 

292 """ 

293 return unpack(value, 8, endian, sign) 

294 

295 

296def u16(value: bytes, endian: str = "little", sign: bool = False) -> int: 

297 """Unpack a 16 bit integer. 

298 

299 Arguments: 

300 value: Value to unpack. 

301 endian: Endianness to use (little, big, network, <, > or !) 

302 sign: Signedness of the integer. 

303 """ 

304 return unpack(value, 16, endian, sign) 

305 

306 

307def u32(value: bytes, endian: str = "little", sign: bool = False) -> int: 

308 """Unpack a 32 bit integer. 

309 

310 Arguments: 

311 value: Value to unpack. 

312 endian: Endianness to use (little, big, network, <, > or !) 

313 sign: Signedness of the integer. 

314 """ 

315 return unpack(value, 32, endian, sign) 

316 

317 

318def u64(value: bytes, endian: str = "little", sign: bool = False) -> int: 

319 """Unpack a 64 bit integer. 

320 

321 Arguments: 

322 value: Value to unpack. 

323 endian: Endianness to use (little, big, network, <, > or !) 

324 sign: Signedness of the integer. 

325 """ 

326 return unpack(value, 64, endian, sign) 

327 

328 

329def swap(value: int, size: int) -> int: 

330 """Swap the endianness of an integer with a given bit size. 

331 

332 Arguments: 

333 value: Integer to swap. 

334 size: Integer size in bits. 

335 """ 

336 return unpack(pack(value, size, ">"), size, "<") 

337 

338 

339def swap16(value: int) -> int: 

340 """Swap the endianness of a 16 bit integer. 

341 

342 Arguments: 

343 value: Integer to swap. 

344 """ 

345 return swap(value, 16) 

346 

347 

348def swap32(value: int) -> int: 

349 """Swap the endianness of a 32 bit integer. 

350 

351 Arguments: 

352 value: Integer to swap. 

353 """ 

354 return swap(value, 32) 

355 

356 

357def swap64(value: int) -> int: 

358 """Swap the endianness of a 64 bit integer. 

359 

360 Arguments: 

361 value: Integer to swap. 

362 """ 

363 return swap(value, 64)