Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/unblob/models.py: 75%

1from __future__ import annotations 

2 

3import abc 

4import dataclasses 

5import itertools 

6import json 

7from enum import Enum 

8from pathlib import Path # noqa: TC003 

9from typing import TYPE_CHECKING, Generic, TypeVar 

10 

11import attrs 

12from pydantic import BaseModel, TypeAdapter, field_validator 

13from structlog import get_logger 

14 

15from .file_utils import Endian, File, InvalidInputFormat, StructParser 

16from .identifiers import new_id 

17from .parser import hexstring2regex 

18from .report import ( 

19 CarveDirectoryReport, 

20 ChunkReport, 

21 EncryptionMetadataReport, 

22 ErrorReport, 

23 MetadataReport, 

24 MultiFileReport, 

25 RandomnessReport, 

26 Report, 

27 UnknownChunkReport, 

28 validate_report_list, 

29) 

30 

31if TYPE_CHECKING: 

32 from collections.abc import Iterable 

33 

34__all__ = [ 

35 "Blob", 

36 "Chunk", 

37 "DExtractor", 

38 "DirectoryExtractor", 

39 "DirectoryHandler", 

40 "DirectoryHandlers", 

41 "DirectoryPattern", 

42 "Endian", 

43 "ExtractError", 

44 "ExtractResult", 

45 "Extractor", 

46 "File", 

47 "Glob", 

48 "Handler", 

49 "HandlerDoc", 

50 "HandlerType", 

51 "Handlers", 

52 "HexString", 

53 "InvalidInputFormat", 

54 "MultiFile", 

55 "PaddingChunk", 

56 "Pattern", 

57 "ProcessResult", 

58 "Reference", 

59 "Regex", 

60 "ReportModel", 

61 "ReportModelAdapter", 

62 "SingleFile", 

63 "StructHandler", 

64 "StructParser", 

65 "TExtractor", 

66 "Task", 

67 "TaskResult", 

68 "UnknownChunk", 

69 "ValidChunk", 

70] 

71 

72logger = get_logger() 

73 

74# The state transitions are: 

75# 

76# file ──► pattern match ──► ValidChunk 

77# 

78 

79 

80class HandlerType(Enum): 

81 ARCHIVE = "Archive" 

82 COMPRESSION = "Compression" 

83 FILESYSTEM = "FileSystem" 

84 EXECUTABLE = "Executable" 

85 BAREMETAL = "Baremetal" 

86 BOOTLOADER = "Bootloader" 

87 ENCRYPTION = "Encryption" 

88 

89 

90@dataclasses.dataclass(frozen=True) 

91class Reference: 

92 title: str 

93 url: str 

94 

95 

96@dataclasses.dataclass 

97class HandlerDoc: 

98 name: str 

99 description: str | None 

100 vendor: str | None 

101 references: list[Reference] 

102 limitations: list[str] 

103 handler_type: HandlerType 

104 private: bool = False 

105 fully_supported: bool = dataclasses.field(init=False) 

106 

107 def __post_init__(self): 

108 self.fully_supported = len(self.limitations) == 0 

109 

110 

111class Task(BaseModel): 

112 path: Path 

113 depth: int 

114 blob_id: str 

115 is_multi_file: bool = False 

116 

117 

118@attrs.define 

119class Blob: 

120 id: str = attrs.field( 

121 factory=new_id, 

122 ) 

123 

124 

125@attrs.define 

126class Chunk(Blob): 

127 """File chunk, have start and end offset, but still can be invalid. 

128 

129 For an array ``b``, a chunk ``c`` represents the slice: 

130 :: 

131 

132 b[c.start_offset:c.end_offset] 

133 """ 

134 

135 start_offset: int = attrs.field(kw_only=True) 

136 """The index of the first byte of the chunk""" 

137 

138 end_offset: int = attrs.field(kw_only=True) 

139 """The index of the first byte after the end of the chunk""" 

140 

141 file: File | None = None 

142 

143 def __attrs_post_init__(self): 

144 if self.start_offset < 0 or self.end_offset < 0: 

145 raise InvalidInputFormat(f"Chunk has negative offset: {self}") 

146 if self.start_offset >= self.end_offset: 

147 raise InvalidInputFormat( 

148 f"Chunk has higher start_offset than end_offset: {self}" 

149 ) 

150 

151 @property 

152 def size(self) -> int: 

153 return self.end_offset - self.start_offset 

154 

155 @property 

156 def range_hex(self) -> str: 

157 return f"0x{self.start_offset:x}-0x{self.end_offset:x}" 

158 

159 @property 

160 def is_whole_file(self): 

161 assert self.file 

162 return self.start_offset == 0 and self.end_offset == self.file.size() 

163 

164 def contains(self, other: Chunk) -> bool: 

165 return ( 

166 self.start_offset < other.start_offset 

167 and self.end_offset >= other.end_offset 

168 ) or ( 

169 self.start_offset <= other.start_offset 

170 and self.end_offset > other.end_offset 

171 ) 

172 

173 def contains_offset(self, offset: int) -> bool: 

174 return self.start_offset <= offset < self.end_offset 

175 

176 def __repr__(self) -> str: 

177 return self.range_hex 

178 

179 

180@attrs.define(repr=False) 

181class ValidChunk(Chunk): 

182 """Known to be valid chunk of a File, can be extracted with an external program.""" 

183 

184 handler: Handler = attrs.field(init=False, eq=False) 

185 metadata_reports: list[MetadataReport] = attrs.field(factory=list, kw_only=True) 

186 

187 @property 

188 def is_encrypted(self) -> bool: 

189 return any( 

190 report.is_encrypted 

191 for report in self.metadata_reports 

192 if isinstance(report, EncryptionMetadataReport) 

193 ) 

194 

195 def extract(self, inpath: Path, outdir: Path) -> ExtractResult | None: 

196 if self.is_encrypted: 

197 logger.warning( 

198 "Encrypted file is not extracted", 

199 path=inpath, 

200 chunk=self, 

201 ) 

202 raise ExtractError 

203 

204 return self.handler.extract(inpath, outdir) 

205 

206 def as_report(self, extraction_reports: list[Report]) -> ChunkReport: 

207 return ChunkReport( 

208 id=self.id, 

209 start_offset=self.start_offset, 

210 end_offset=self.end_offset, 

211 size=self.size, 

212 handler_name=self.handler.NAME, 

213 metadata_reports=self.metadata_reports, 

214 extraction_reports=extraction_reports, 

215 ) 

216 

217 

218@attrs.define(repr=False) 

219class UnknownChunk(Chunk): 

220 r"""Gaps between valid chunks or otherwise unknown chunks. 

221 

222 Important for manual analysis, and analytical certainty: for example 

223 randomness, other chunks inside it, metadata, etc. 

224 

225 These are not extracted, just logged for information purposes and further analysis, 

226 like most common bytes (like \x00 and \xFF), ASCII strings, high randomness, etc. 

227 """ 

228 

229 def as_report(self, randomness: RandomnessReport | None) -> UnknownChunkReport: 

230 return UnknownChunkReport( 

231 id=self.id, 

232 start_offset=self.start_offset, 

233 end_offset=self.end_offset, 

234 size=self.size, 

235 randomness=randomness, 

236 ) 

237 

238 

239@attrs.define(repr=False) 

240class PaddingChunk(Chunk): 

241 r"""Gaps between valid chunks or otherwise unknown chunks. 

242 

243 Important for manual analysis, and analytical certanity: for example 

244 randomness, other chunks inside it, metadata, etc. 

245 """ 

246 

247 def as_report( 

248 self, 

249 randomness: RandomnessReport | None, # noqa: ARG002 

250 ) -> ChunkReport: 

251 return ChunkReport( 

252 id=self.id, 

253 start_offset=self.start_offset, 

254 end_offset=self.end_offset, 

255 size=self.size, 

256 handler_name="padding", 

257 extraction_reports=[], 

258 metadata_reports=[], 

259 ) 

260 

261 

262@attrs.define 

263class MultiFile(Blob): 

264 name: str = attrs.field(kw_only=True) 

265 paths: list[Path] = attrs.field(kw_only=True) 

266 

267 handler: DirectoryHandler = attrs.field(init=False, eq=False) 

268 

269 def extract(self, outdir: Path) -> ExtractResult | None: 

270 return self.handler.extract(self.paths, outdir) 

271 

272 def as_report(self, extraction_reports: list[Report]) -> MultiFileReport: 

273 return MultiFileReport( 

274 id=self.id, 

275 name=self.name, 

276 paths=self.paths, 

277 handler_name=self.handler.NAME, 

278 extraction_reports=extraction_reports, 

279 metadata_reports=[], 

280 ) 

281 

282 

283ReportType = TypeVar("ReportType", bound=Report) 

284 

285 

286class TaskResult(BaseModel): 

287 task: Task 

288 reports: list[Report] = [] 

289 subtasks: list[Task] = [] 

290 

291 @field_validator("reports", mode="before") 

292 @classmethod 

293 def validate_reports(cls, value): 

294 return validate_report_list(value) 

295 

296 def add_report(self, report: Report): 

297 self.reports.append(report) 

298 

299 def add_subtask(self, task: Task): 

300 self.subtasks.append(task) 

301 

302 def filter_reports(self, report_class: type[ReportType]) -> list[ReportType]: 

303 return [report for report in self.reports if isinstance(report, report_class)] 

304 

305 

306class ProcessResult(BaseModel): 

307 results: list[TaskResult] = [] 

308 

309 @property 

310 def errors(self) -> list[ErrorReport]: 

311 reports = itertools.chain.from_iterable(r.reports for r in self.results) 

312 interesting_reports = ( 

313 r for r in reports if isinstance(r, ErrorReport | ChunkReport) 

314 ) 

315 errors = [] 

316 for report in interesting_reports: 

317 if isinstance(report, ErrorReport): 

318 errors.append(report) 

319 else: 

320 errors.extend( 

321 r for r in report.extraction_reports if isinstance(r, ErrorReport) 

322 ) 

323 return errors 

324 

325 def register(self, result: TaskResult): 

326 self.results.append(result) 

327 

328 def to_json(self, indent=" "): 

329 return json.dumps( 

330 [ 

331 result.model_dump(mode="json", serialize_as_any=True) 

332 for result in self.results 

333 ], 

334 indent=indent, 

335 ) 

336 

337 def get_output_dir(self) -> Path | None: 

338 try: 

339 top_result = self.results[0] 

340 if carves := top_result.filter_reports(CarveDirectoryReport): 

341 # we have a top level carve 

342 return carves[0].carve_dir 

343 

344 # we either have an extraction, 

345 # and the extract directory registered as subtask 

346 return top_result.subtasks[0].path 

347 except IndexError: 

348 # or no extraction 

349 return None 

350 

351 

352ReportModel = list[TaskResult] 

353ReportModelAdapter = TypeAdapter(ReportModel) 

354"""Use this for deserialization (import JSON report back into Python 

355objects) of the JSON report. 

356 

357For example: 

358 

359with open('report.json', 'r') as f: 

360 data = f.read() 

361 report_data = ReportModelAdapter.validate_json(data) 

362 

363For another example see: 

364tests/test_models.py::Test_to_json::test_process_result_deserialization 

365""" 

366 

367 

368class ExtractError(Exception): 

369 """There was an error during extraction.""" 

370 

371 def __init__(self, *reports: Report): 

372 super().__init__() 

373 self.reports: tuple[Report, ...] = reports 

374 

375 

376@attrs.define(kw_only=True) 

377class ExtractResult: 

378 reports: list[Report] 

379 

380 

381class Extractor(abc.ABC): 

382 def get_dependencies(self) -> list[str]: 

383 """Return the external command dependencies.""" 

384 return [] 

385 

386 @abc.abstractmethod 

387 def extract(self, inpath: Path, outdir: Path) -> ExtractResult | None: 

388 """Extract the carved out chunk. 

389 

390 Raises ExtractError on failure. 

391 """ 

392 

393 

394class DirectoryExtractor(abc.ABC): 

395 def get_dependencies(self) -> list[str]: 

396 """Return the external command dependencies.""" 

397 return [] 

398 

399 @abc.abstractmethod 

400 def extract(self, paths: list[Path], outdir: Path) -> ExtractResult | None: 

401 """Extract from a multi file path list. 

402 

403 Raises ExtractError on failure. 

404 """ 

405 

406 

407class Pattern(str): 

408 def as_regex(self) -> bytes: 

409 raise NotImplementedError 

410 

411 

412class HexString(Pattern): 

413 """Hex string can be a YARA rule like hexadecimal string. 

414 

415 It is useful to simplify defining binary strings using hex 

416 encoding, wild-cards, jumps and alternatives. Hexstrings are 

417 convereted to hyperscan compatible PCRE regex. 

418 

419 See YARA & Hyperscan documentation for more details: 

420 

421 - https://yara.readthedocs.io/en/stable/writingrules.html#hexadecimal-strings 

422 

423 - https://intel.github.io/hyperscan/dev-reference/compilation.html#pattern-support 

424 

425 You can specify the following: 

426 

427 - normal bytes using hexadecimals: 01 de ad co de ff 

428 

429 - wild-cards can match single bytes and can be mixed with 

430 normal hex: 01 ?? 02 

431 

432 - wild-cards can also match first and second nibles: 0? ?0 

433 

434 - jumps can be specified for multiple wildcard bytes: [3] 

435 [2-5] 

436 

437 - alternatives can be specified as well: ( 01 02 | 03 04 ) The 

438 above can be combined and alternatives nested: 01 02 ( 03 04 

439 | (0? | 03 | ?0) | 05 ?? ) 06 

440 

441 Single line comments can be specified using // 

442 

443 We do NOT support the following YARA syntax: 

444 

445 - comments using /* */ notation 

446 

447 - infinite jumps: [-] 

448 

449 - unbounded jumps: [3-] or [-4] (use [0-4] instead) 

450 """ 

451 

452 def as_regex(self) -> bytes: 

453 return hexstring2regex(self) 

454 

455 

456class Regex(Pattern): 

457 """Byte PCRE regex. 

458 

459 See hyperscan documentation for more details: 

460 https://intel.github.io/hyperscan/dev-reference/compilation.html#pattern-support. 

461 """ 

462 

463 def as_regex(self) -> bytes: 

464 return self.encode() 

465 

466 

467class DirectoryPattern: 

468 def get_files(self, directory: Path) -> Iterable[Path]: 

469 raise NotImplementedError 

470 

471 

472class Glob(DirectoryPattern): 

473 def __init__(self, *patterns): 

474 if not patterns: 

475 raise ValueError("At least one pattern must be provided") 

476 self._patterns = patterns 

477 

478 def get_files(self, directory: Path) -> Iterable[Path]: 

479 for pattern in self._patterns: 

480 yield from directory.glob(pattern) 

481 

482 

483class SingleFile(DirectoryPattern): 

484 def __init__(self, filename): 

485 self._filename = filename 

486 

487 def get_files(self, directory: Path) -> Iterable[Path]: 

488 path = directory / self._filename 

489 return [path] if path.exists() else [] 

490 

491 

492DExtractor = TypeVar("DExtractor", bound=DirectoryExtractor | None) 

493 

494 

495class DirectoryHandler(abc.ABC, Generic[DExtractor]): 

496 """A directory type handler is responsible for searching, validating and "unblobbing" files from multiple files in a directory.""" 

497 

498 NAME: str 

499 

500 EXTRACTOR: DExtractor 

501 

502 PATTERN: DirectoryPattern 

503 

504 DOC: HandlerDoc | None 

505 

506 @classmethod 

507 def get_dependencies(cls): 

508 """Return external command dependencies needed for this handler to work.""" 

509 if cls.EXTRACTOR is not None: 

510 return cls.EXTRACTOR.get_dependencies() 

511 return [] 

512 

513 @abc.abstractmethod 

514 def calculate_multifile(self, file: Path) -> MultiFile | None: 

515 """Calculate the MultiFile in a directory, using a file matched by the pattern as a starting point.""" 

516 

517 def extract(self, paths: list[Path], outdir: Path) -> ExtractResult | None: 

518 if self.EXTRACTOR is None: 

519 logger.debug("Skipping file: no extractor.", paths=paths) 

520 raise ExtractError 

521 

522 # We only extract every blob once, it's a mistake to extract the same blob again 

523 outdir.mkdir(parents=True, exist_ok=False) 

524 

525 return self.EXTRACTOR.extract(paths, outdir) 

526 

527 

528TExtractor = TypeVar("TExtractor", bound=Extractor | None) 

529 

530 

531class Handler(abc.ABC, Generic[TExtractor]): 

532 """A file type handler is responsible for searching, validating and "unblobbing" files from Blobs.""" 

533 

534 NAME: str 

535 PATTERNS: list[Pattern] 

536 # We need this, because not every match reflects the actual start 

537 # (e.g. tar magic is in the middle of the header) 

538 PATTERN_MATCH_OFFSET: int = 0 

539 

540 EXTRACTOR: TExtractor 

541 

542 DOC: HandlerDoc | None 

543 

544 @classmethod 

545 def get_dependencies(cls): 

546 """Return external command dependencies needed for this handler to work.""" 

547 if cls.EXTRACTOR is not None: 

548 return cls.EXTRACTOR.get_dependencies() 

549 return [] 

550 

551 @abc.abstractmethod 

552 def calculate_chunk(self, file: File, start_offset: int) -> ValidChunk | None: 

553 """Calculate the Chunk offsets from the File and the file type headers.""" 

554 

555 def extract(self, inpath: Path, outdir: Path) -> ExtractResult | None: 

556 if self.EXTRACTOR is None: 

557 logger.debug("Skipping file: no extractor.", path=inpath) 

558 raise ExtractError 

559 

560 # We only extract every blob once, it's a mistake to extract the same blob again 

561 outdir.mkdir(parents=True, exist_ok=False) 

562 

563 return self.EXTRACTOR.extract(inpath, outdir) 

564 

565 

566class StructHandler(Handler): 

567 C_DEFINITIONS: str 

568 # A struct from the C_DEFINITIONS used to parse the file's header 

569 HEADER_STRUCT: str 

570 

571 def __init__(self): 

572 self._struct_parser = StructParser(self.C_DEFINITIONS) 

573 

574 @property 

575 def cparser_le(self): 

576 return self._struct_parser.cparser_le 

577 

578 @property 

579 def cparser_be(self): 

580 return self._struct_parser.cparser_be 

581 

582 def parse_header(self, file: File, endian=Endian.LITTLE): 

583 header = self._struct_parser.parse(self.HEADER_STRUCT, file, endian) 

584 logger.debug("Header parsed", header=header, _verbosity=3) 

585 return header 

586 

587 

588Handlers = tuple[type[Handler], ...] 

589DirectoryHandlers = tuple[type[DirectoryHandler], ...]