wheelfile.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241
  1. from __future__ import annotations
  2. __all__ = ["WHEEL_INFO_RE", "WheelFile", "WheelError"]
  3. import base64
  4. import csv
  5. import hashlib
  6. import logging
  7. import os.path
  8. import re
  9. import stat
  10. import time
  11. from io import StringIO, TextIOWrapper
  12. from typing import IO, TYPE_CHECKING, Literal
  13. from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo
  14. if TYPE_CHECKING:
  15. from _typeshed import SizedBuffer, StrPath
  16. # Non-greedy matching of an optional build number may be too clever (more
  17. # invalid wheel filenames will match). Separate regex for .dist-info?
  18. WHEEL_INFO_RE = re.compile(
  19. r"""^(?P<namever>(?P<name>[^\s-]+?)-(?P<ver>[^\s-]+?))(-(?P<build>\d[^\s-]*))?
  20. -(?P<pyver>[^\s-]+?)-(?P<abi>[^\s-]+?)-(?P<plat>\S+)\.whl$""",
  21. re.VERBOSE,
  22. )
  23. MINIMUM_TIMESTAMP = 315532800 # 1980-01-01 00:00:00 UTC
  24. log = logging.getLogger("wheel")
  25. class WheelError(Exception):
  26. pass
  27. def urlsafe_b64encode(data: bytes) -> bytes:
  28. """urlsafe_b64encode without padding"""
  29. return base64.urlsafe_b64encode(data).rstrip(b"=")
  30. def urlsafe_b64decode(data: bytes) -> bytes:
  31. """urlsafe_b64decode without padding"""
  32. pad = b"=" * (4 - (len(data) & 3))
  33. return base64.urlsafe_b64decode(data + pad)
  34. def get_zipinfo_datetime(
  35. timestamp: float | None = None,
  36. ) -> tuple[int, int, int, int, int]:
  37. # Some applications need reproducible .whl files, but they can't do this without
  38. # forcing the timestamp of the individual ZipInfo objects. See issue #143.
  39. timestamp = int(os.environ.get("SOURCE_DATE_EPOCH", timestamp or time.time()))
  40. timestamp = max(timestamp, MINIMUM_TIMESTAMP)
  41. return time.gmtime(timestamp)[0:6]
  42. class WheelFile(ZipFile):
  43. """A ZipFile derivative class that also reads SHA-256 hashes from
  44. .dist-info/RECORD and checks any read files against those.
  45. """
  46. _default_algorithm = hashlib.sha256
  47. def __init__(
  48. self,
  49. file: StrPath,
  50. mode: Literal["r", "w", "x", "a"] = "r",
  51. compression: int = ZIP_DEFLATED,
  52. ):
  53. basename = os.path.basename(file)
  54. self.parsed_filename = WHEEL_INFO_RE.match(basename)
  55. if not basename.endswith(".whl") or self.parsed_filename is None:
  56. raise WheelError(f"Bad wheel filename {basename!r}")
  57. ZipFile.__init__(self, file, mode, compression=compression, allowZip64=True)
  58. self.dist_info_path = "{}.dist-info".format(
  59. self.parsed_filename.group("namever")
  60. )
  61. self.record_path = self.dist_info_path + "/RECORD"
  62. self._file_hashes: dict[str, tuple[None, None] | tuple[int, bytes]] = {}
  63. self._file_sizes = {}
  64. if mode == "r":
  65. # Ignore RECORD and any embedded wheel signatures
  66. self._file_hashes[self.record_path] = None, None
  67. self._file_hashes[self.record_path + ".jws"] = None, None
  68. self._file_hashes[self.record_path + ".p7s"] = None, None
  69. # Fill in the expected hashes by reading them from RECORD
  70. try:
  71. record = self.open(self.record_path)
  72. except KeyError:
  73. raise WheelError(f"Missing {self.record_path} file") from None
  74. with record:
  75. for line in csv.reader(
  76. TextIOWrapper(record, newline="", encoding="utf-8")
  77. ):
  78. path, hash_sum, size = line
  79. if not hash_sum:
  80. continue
  81. algorithm, hash_sum = hash_sum.split("=")
  82. try:
  83. hashlib.new(algorithm)
  84. except ValueError:
  85. raise WheelError(
  86. f"Unsupported hash algorithm: {algorithm}"
  87. ) from None
  88. if algorithm.lower() in {"md5", "sha1"}:
  89. raise WheelError(
  90. f"Weak hash algorithm ({algorithm}) is not permitted by "
  91. f"PEP 427"
  92. )
  93. self._file_hashes[path] = (
  94. algorithm,
  95. urlsafe_b64decode(hash_sum.encode("ascii")),
  96. )
  97. def open(
  98. self,
  99. name_or_info: str | ZipInfo,
  100. mode: Literal["r", "w"] = "r",
  101. pwd: bytes | None = None,
  102. ) -> IO[bytes]:
  103. def _update_crc(newdata: bytes) -> None:
  104. eof = ef._eof
  105. update_crc_orig(newdata)
  106. running_hash.update(newdata)
  107. if eof and running_hash.digest() != expected_hash:
  108. raise WheelError(f"Hash mismatch for file '{ef_name}'")
  109. ef_name = (
  110. name_or_info.filename if isinstance(name_or_info, ZipInfo) else name_or_info
  111. )
  112. if (
  113. mode == "r"
  114. and not ef_name.endswith("/")
  115. and ef_name not in self._file_hashes
  116. ):
  117. raise WheelError(f"No hash found for file '{ef_name}'")
  118. ef = ZipFile.open(self, name_or_info, mode, pwd)
  119. if mode == "r" and not ef_name.endswith("/"):
  120. algorithm, expected_hash = self._file_hashes[ef_name]
  121. if expected_hash is not None:
  122. # Monkey patch the _update_crc method to also check for the hash from
  123. # RECORD
  124. running_hash = hashlib.new(algorithm)
  125. update_crc_orig, ef._update_crc = ef._update_crc, _update_crc
  126. return ef
  127. def write_files(self, base_dir: str) -> None:
  128. log.info("creating %r and adding %r to it", self.filename, base_dir)
  129. deferred: list[tuple[str, str]] = []
  130. for root, dirnames, filenames in os.walk(base_dir):
  131. # Sort the directory names so that `os.walk` will walk them in a
  132. # defined order on the next iteration.
  133. dirnames.sort()
  134. for name in sorted(filenames):
  135. path = os.path.normpath(os.path.join(root, name))
  136. if os.path.isfile(path):
  137. arcname = os.path.relpath(path, base_dir).replace(os.path.sep, "/")
  138. if arcname == self.record_path:
  139. pass
  140. elif root.endswith(".dist-info"):
  141. deferred.append((path, arcname))
  142. else:
  143. self.write(path, arcname)
  144. deferred.sort()
  145. for path, arcname in deferred:
  146. self.write(path, arcname)
  147. def write(
  148. self,
  149. filename: str,
  150. arcname: str | None = None,
  151. compress_type: int | None = None,
  152. ) -> None:
  153. with open(filename, "rb") as f:
  154. st = os.fstat(f.fileno())
  155. data = f.read()
  156. zinfo = ZipInfo(
  157. arcname or filename, date_time=get_zipinfo_datetime(st.st_mtime)
  158. )
  159. zinfo.external_attr = (stat.S_IMODE(st.st_mode) | stat.S_IFMT(st.st_mode)) << 16
  160. zinfo.compress_type = compress_type or self.compression
  161. self.writestr(zinfo, data, compress_type)
  162. def writestr(
  163. self,
  164. zinfo_or_arcname: str | ZipInfo,
  165. data: SizedBuffer | str,
  166. compress_type: int | None = None,
  167. ) -> None:
  168. if isinstance(zinfo_or_arcname, str):
  169. zinfo_or_arcname = ZipInfo(
  170. zinfo_or_arcname, date_time=get_zipinfo_datetime()
  171. )
  172. zinfo_or_arcname.compress_type = self.compression
  173. zinfo_or_arcname.external_attr = (0o664 | stat.S_IFREG) << 16
  174. if isinstance(data, str):
  175. data = data.encode("utf-8")
  176. ZipFile.writestr(self, zinfo_or_arcname, data, compress_type)
  177. fname = (
  178. zinfo_or_arcname.filename
  179. if isinstance(zinfo_or_arcname, ZipInfo)
  180. else zinfo_or_arcname
  181. )
  182. log.info("adding %r", fname)
  183. if fname != self.record_path:
  184. hash_ = self._default_algorithm(data)
  185. self._file_hashes[fname] = (
  186. hash_.name,
  187. urlsafe_b64encode(hash_.digest()).decode("ascii"),
  188. )
  189. self._file_sizes[fname] = len(data)
  190. def close(self) -> None:
  191. # Write RECORD
  192. if self.fp is not None and self.mode == "w" and self._file_hashes:
  193. data = StringIO()
  194. writer = csv.writer(data, delimiter=",", quotechar='"', lineterminator="\n")
  195. writer.writerows(
  196. (
  197. (fname, algorithm + "=" + hash_, self._file_sizes[fname])
  198. for fname, (algorithm, hash_) in self._file_hashes.items()
  199. )
  200. )
  201. writer.writerow((format(self.record_path), "", ""))
  202. self.writestr(self.record_path, data.getvalue())
  203. ZipFile.close(self)