Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/model_signing/manifest.py: 92%

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

146 statements  

1# Copyright 2024 The Sigstore Authors 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""An in-memory serialized representation of an ML model. 

16 

17A manifest pairs objects from the model (e.g., files, shards of files), with 

18their hashes. When signing the model we first generate a manifest from 

19serializing the model using a configured serialization method (see 

20`model_signing.signing`). 

21 

22When verifying the integrity of the model, after checking the authenticity of 

23the signature, we extract a manifest from it. Then, we serialize the local model 

24(the model being tested) and compare the two manifests. 

25 

26The serialization method used during signing must match the one used during 

27verification. We can auto detect the method to use during verification from the 

28signature, but it is recommended to be explicit when possible. 

29 

30Comparing the manifests can be done by checking that every item matches, both in 

31name and in associated hash. In the future we will support partial object 

32matching. This is useful, for example, for the cases where the original model 

33contained files for multiple ML frameworks, but the user only uses the model 

34with one framework. This way, the user can verify the integrity only for the 

35files that are actually used. 

36 

37This API should not be used directly, we don't guarantee that it is fully stable 

38at the moment. 

39""" 

40 

41import abc 

42from collections.abc import Iterable, Iterator 

43import dataclasses 

44import pathlib 

45from typing import Any, Final, Self 

46 

47from typing_extensions import override 

48 

49from model_signing._hashing import hashing 

50 

51 

52@dataclasses.dataclass(frozen=True, order=True) 

53class _ResourceDescriptor: 

54 """A description of any content from any `Manifest`. 

55 

56 We aim this to be similar to in-toto's `ResourceDescriptor`. To support 

57 cases where in-toto cannot be directly used, we make this a dataclass that 

58 can be mapped to in-toto when needed, and used as its own otherwise. 

59 

60 Not all fields from in-toto are specified at this moment. All fields here 

61 must be present, unlike in-toto, where all are optional. 

62 

63 See github.com/in-toto/attestation/blob/main/spec/v1/resource_descriptor.md 

64 for the in-toto specification. 

65 

66 Attributes: 

67 identifier: A string that uniquely identifies this object within the 

68 manifest. Depending on serialized format, users might require the 

69 identifier to be unique across all manifests stored in a system. 

70 Producers and consumers can agree on additional requirements (e.g., 

71 several descriptors must have a common pattern for the identifier and 

72 the integrity of the model implies integrity of all these items, 

73 ignoring any other descriptor). Corresponds to `name`, `uri`, or 

74 `content` in in-toto specification. 

75 digest: One digest for the item. Note that unlike in-toto, we only have 

76 one digest for the item and it is always required. 

77 """ 

78 

79 identifier: str 

80 digest: hashing.Digest 

81 

82 

83class _ManifestKey(metaclass=abc.ABCMeta): 

84 """An object that can be a key for an item in a manifest. 

85 

86 We need to be able to convert the key to string when we serialize the 

87 manifest and to rebuild the object from the serialized representation when 

88 we deserialize the manifest from the signature. 

89 """ 

90 

91 @classmethod 

92 @abc.abstractmethod 

93 def from_str(cls, s: str) -> Self: 

94 """Builds a manifest key from the string representation. 

95 

96 It is guaranteed that for any `key` of a derived class `C` and a (valid) 

97 string `s`, the following two round-trip properties hold: 

98 

99 ``` 

100 str(C.from_str(s)) == s 

101 C.from_str(str(key)) == key 

102 ``` 

103 

104 Raises: 

105 ValueError: if the string argument cannot be decoded correctly. 

106 """ 

107 

108 

109@dataclasses.dataclass(frozen=True, order=True) 

110class _File(_ManifestKey): 

111 """A dataclass to hold information about a file as a manifest key. 

112 

113 Attributes: 

114 path: The path to the file, relative to the model root. 

115 """ 

116 

117 path: pathlib.PurePath 

118 

119 def __str__(self) -> str: 

120 return str(self.path) 

121 

122 @classmethod 

123 @override 

124 def from_str(cls, s: str) -> Self: 

125 # Note that we always decode the string to a pure POSIX path 

126 return cls(pathlib.PurePosixPath(s)) 

127 

128 

129@dataclasses.dataclass(frozen=True, order=True) 

130class _Shard(_ManifestKey): 

131 """A dataclass to hold information about a file shard as a manifest key. 

132 

133 Attributes: 

134 path: The path to the file, relative to the model root. 

135 start: The start offset of the shard (included). 

136 end: The end offset of the shard (not included). 

137 """ 

138 

139 path: pathlib.PurePath 

140 start: int 

141 end: int 

142 

143 def __str__(self) -> str: 

144 return f"{str(self.path)}:{self.start}:{self.end}" 

145 

146 @classmethod 

147 @override 

148 def from_str(cls, s: str) -> Self: 

149 parts = s.split(":") 

150 if len(parts) != 3: 

151 raise ValueError(f"Expected 3 components separated by `:`, got {s}") 

152 

153 path = pathlib.PurePosixPath(parts[0]) 

154 start = int(parts[1]) 

155 end = int(parts[2]) 

156 

157 return cls(path, start, end) 

158 

159 

160class ManifestItem(metaclass=abc.ABCMeta): 

161 """An individual object of a model, stored as an item in a manifest. 

162 

163 For example, this could be a file, or a file shard. All file paths are 

164 relative to the model root, to allow moving or renaming the model, without 

165 invalidating the signature. 

166 

167 The integrity of each `ManifestItem` can be verified individually and in 

168 parallel. If the item is backed by a file, we recompute the hash of the 

169 portion of the file that represents this item. 

170 

171 Attributes: 

172 digest: The digest of the item. Use the `key` property to obtain a 

173 canonical unique representation for the item. 

174 """ 

175 

176 digest: hashing.Digest 

177 

178 @property 

179 @abc.abstractmethod 

180 def key(self) -> _ManifestKey: 

181 """A unique representation for the manifest item. 

182 

183 Two items in the same manifest must not share the same `key`. The 

184 information contained in `key` should be sufficient to determine how to 

185 compute the item's digest. 

186 """ 

187 

188 

189class FileManifestItem(ManifestItem): 

190 """A manifest item that records a filename path together with its digest. 

191 

192 Note that the path component is a `pathlib.PurePath`, relative to the model. 

193 To ensure that the manifest is consistent across operating systems, we 

194 convert the path to a POSIX path. 

195 """ 

196 

197 def __init__(self, *, path: pathlib.PurePath, digest: hashing.Digest): 

198 """Builds a manifest item pairing a file with its digest. 

199 

200 Args: 

201 path: The path to the file, relative to the model root. 

202 digest: The digest of the file. 

203 """ 

204 # Note: we need to force a PurePosixPath to canonicalize the manifest. 

205 self._path = pathlib.PurePosixPath(path) 

206 self.digest = digest 

207 

208 @property 

209 @override 

210 def key(self) -> _ManifestKey: 

211 return _File(self._path) 

212 

213 

214class ShardedFileManifestItem(ManifestItem): 

215 """A manifest item that records a file shard together with its digest. 

216 

217 Note that the path component is a `pathlib.PurePath`, relative to the model. 

218 To ensure that the manifest is consistent across operating systems, we 

219 convert the path to a POSIX path. 

220 """ 

221 

222 def __init__( 

223 self, 

224 *, 

225 path: pathlib.PurePath, 

226 start: int, 

227 end: int, 

228 digest: hashing.Digest, 

229 ): 

230 """Builds a manifest item pairing a file shard with its digest. 

231 

232 Args: 

233 path: The path to the file, relative to the model root. 

234 start: The start offset of the shard (included). 

235 end: The end offset of the shard (not included). 

236 digest: The digest of the file shard. 

237 """ 

238 # Note: we need to force a PurePosixPath to canonicalize the manifest. 

239 self._path = pathlib.PurePosixPath(path) 

240 self._start = start 

241 self._end = end 

242 self.digest = digest 

243 

244 @property 

245 @override 

246 def key(self) -> _ManifestKey: 

247 return _Shard(self._path, self._start, self._end) 

248 

249 

250class SerializationType(metaclass=abc.ABCMeta): 

251 """A description of the serialization process that generated the manifest. 

252 

253 These should record all the parameters needed to ensure a reproducible 

254 serialization. These are used to build a manifest back from the signature in 

255 a backwards compatible way. We use these to determine what serialization to 

256 use when verifying a signature. 

257 """ 

258 

259 @property 

260 @abc.abstractmethod 

261 def serialization_parameters(self) -> dict[str, Any]: 

262 """The arguments of the serialization method.""" 

263 

264 @classmethod 

265 def from_args(cls, args: dict[str, Any]) -> Self: 

266 """Builds an instance of this class based on the dict representation. 

267 

268 This is the reverse of `serialization_parameters`. 

269 

270 Args: 

271 args: The arguments as a dictionary (equivalent to `**kwargs`). 

272 """ 

273 serialization_type = args["method"] 

274 for subclass in [_FileSerialization, _ShardSerialization]: 

275 if serialization_type == subclass.method: 

276 return subclass._from_args(args) 

277 raise ValueError(f"Unknown serialization type {serialization_type}") 

278 

279 @classmethod 

280 @abc.abstractmethod 

281 def _from_args(cls, args: dict[str, Any]) -> Self: 

282 """Performs the actual build from `from_dict`.""" 

283 

284 @abc.abstractmethod 

285 def new_item(self, name: str, digest: hashing.Digest) -> ManifestItem: 

286 """Builds a `ManifestItem` of the correct type. 

287 

288 Each serialization type results in different types for the items in the 

289 manifest. This method parses the `name` of the item according to the 

290 serialization type to create the proper manifest item. 

291 

292 Args: 

293 name: The name of the item, as shown in the manifest. 

294 digest: The digest of the item. 

295 """ 

296 

297 

298class _FileSerialization(SerializationType): 

299 method: Final[str] = "files" 

300 

301 def __init__( 

302 self, 

303 hash_type: str, 

304 allow_symlinks: bool = False, 

305 ignore_paths: Iterable[pathlib.Path] = frozenset(), 

306 ): 

307 """Records the manifest serialization type for serialization by files. 

308 

309 We only need to record the hashing engine used and whether symlinks are 

310 hashed or ignored. 

311 

312 Args: 

313 hash_type: A string representation of the hash algorithm. 

314 allow_symlinks: Controls whether symbolic links are included. 

315 ignore_paths: Paths of files to ignore. 

316 """ 

317 self._hash_type = hash_type 

318 self._allow_symlinks = allow_symlinks 

319 self._ignore_paths = [str(p) for p in ignore_paths] 

320 

321 @property 

322 @override 

323 def serialization_parameters(self) -> dict[str, Any]: 

324 res = { 

325 "method": self.method, 

326 "hash_type": self._hash_type, 

327 "allow_symlinks": self._allow_symlinks, 

328 } 

329 if self._ignore_paths: 

330 res["ignore_paths"] = self._ignore_paths 

331 return res 

332 

333 @classmethod 

334 @override 

335 def _from_args(cls, args: dict[str, Any]) -> Self: 

336 return cls( 

337 args["hash_type"], 

338 args["allow_symlinks"], 

339 args.get("ignore_paths", frozenset()), 

340 ) 

341 

342 @override 

343 def new_item(self, name: str, digest: hashing.Digest) -> ManifestItem: 

344 path = pathlib.PurePosixPath(name) 

345 return FileManifestItem(path=path, digest=digest) 

346 

347 

348class _ShardSerialization(SerializationType): 

349 method: Final[str] = "shards" 

350 

351 def __init__( 

352 self, 

353 hash_type: str, 

354 shard_size: int, 

355 allow_symlinks: bool = False, 

356 ignore_paths: Iterable[pathlib.Path] = frozenset(), 

357 ): 

358 """Records the manifest serialization type for serialization by shards. 

359 

360 We need to record the hashing engine used and whether symlinks are 

361 hashed or ignored, just like for file serialization. We also need to 

362 record the shard size used to split the files, since different shard 

363 sizes results in different resources. 

364 

365 Args: 

366 hash_type: A string representation of the hash algorithm. 

367 shard_size: The size of each file shard. 

368 allow_symlinks: Controls whether symbolic links are included. 

369 ignore_paths: Paths of files to ignore. 

370 """ 

371 self._hash_type = hash_type 

372 self._allow_symlinks = allow_symlinks 

373 self._shard_size = shard_size 

374 self._ignore_paths = [str(p) for p in ignore_paths] 

375 

376 @property 

377 @override 

378 def serialization_parameters(self) -> dict[str, Any]: 

379 res = { 

380 "method": self.method, 

381 "hash_type": self._hash_type, 

382 "shard_size": self._shard_size, 

383 "allow_symlinks": self._allow_symlinks, 

384 } 

385 if self._ignore_paths: 

386 res["ignore_paths"] = self._ignore_paths 

387 return res 

388 

389 @classmethod 

390 @override 

391 def _from_args(cls, args: dict[str, Any]) -> Self: 

392 return cls( 

393 args["hash_type"], 

394 args["shard_size"], 

395 args["allow_symlinks"], 

396 args.get("ignore_paths", frozenset()), 

397 ) 

398 

399 @override 

400 def new_item(self, name: str, digest: hashing.Digest) -> ManifestItem: 

401 parts = name.split(":") 

402 if len(parts) != 3: 

403 raise ValueError( 

404 "Invalid resource name: expected 3 components separated by " 

405 f"`:`, got {name}" 

406 ) 

407 

408 path = pathlib.PurePosixPath(parts[0]) 

409 start = int(parts[1]) 

410 end = int(parts[2]) 

411 return ShardedFileManifestItem( 

412 path=path, start=start, end=end, digest=digest 

413 ) 

414 

415 

416class Manifest: 

417 """Generic manifest file to represent a model.""" 

418 

419 def __init__( 

420 self, 

421 model_name: str, 

422 items: Iterable[ManifestItem], 

423 serialization_type: SerializationType, 

424 ): 

425 """Builds a manifest from a collection of already hashed objects. 

426 

427 Args: 

428 model_name: A name for the model that generated the manifest. This 

429 is the final component of the model path, and is only informative. 

430 See `model_name` property. 

431 items: An iterable sequence of objects and their hashes. 

432 serialization_type: The serialization method used to create items. 

433 """ 

434 self._name = model_name 

435 self._item_to_digest = {item.key: item.digest for item in items} 

436 self._serialization_type = serialization_type 

437 

438 def __eq__(self, other: Self): 

439 """Check equality based on item digests, ignoring model name.""" 

440 return self._item_to_digest == other._item_to_digest 

441 

442 def resource_descriptors(self) -> Iterator[_ResourceDescriptor]: 

443 """Yields each resource from the manifest, one by one.""" 

444 for item, digest in sorted(self._item_to_digest.items()): 

445 yield _ResourceDescriptor(identifier=str(item), digest=digest) 

446 

447 @property 

448 def model_name(self) -> str: 

449 """The name of the model when serialized (final component of the path). 

450 

451 This is only informative. Changing the name of the model should still 

452 result in the same digests after serialization, it must not invalidate 

453 signatures. As a result, two manifests with different model names but 

454 with the same resource descriptors will compare equal. 

455 """ 

456 return self._name 

457 

458 @property 

459 def serialization_type(self) -> dict[str, Any]: 

460 """The serialization (and arguments) used to build the manifest. 

461 

462 This is needed to record the serialization method used to generate the 

463 manifest so that signature verification can use the same method. 

464 """ 

465 return self._serialization_type.serialization_parameters