Read SBOM inside embed()

Author: radarhere

.github/embed-sbom.py

@@ -22,7 +22,7 @@ def record_entry(path: str, data: bytes) -> str:
     return f"{path},sha256={digest.rstrip(b'=').decode()},{len(data)}"
 
 
-def embed(wheel: Path, sbom_name: str, sbom_bytes: bytes) -> None:
+def embed(wheel: Path, sbom: Path) -> None:
     with zipfile.ZipFile(wheel) as zf:
         infos = zf.infolist()
         contents = {info.filename: zf.read(info.filename) for info in infos}
@@ -33,7 +33,9 @@ def embed(wheel: Path, sbom_name: str, sbom_bytes: bytes) -> None:
         if name.endswith(".dist-info/RECORD") and name.count("/") == 1
     )
     dist_info = record_name.rsplit("/", 1)[0]
-    sbom_path = f"{dist_info}/sboms/{sbom_name}"
+
+    sbom_bytes = sbom.read_bytes()
+    sbom_path = f"{dist_info}/sboms/{sbom.name}"
 
     # Append a matching RECORD line for the SBOM (RECORD's own line has no hash).
     lines = contents[record_name].decode("utf-8").splitlines()
@@ -49,16 +51,18 @@ def embed(wheel: Path, sbom_name: str, sbom_bytes: bytes) -> None:
         zf.writestr(sbom_path, sbom_bytes)
     tmp.replace(wheel)
 
+    print(f"Embedded {sbom.name} in {wheel.name}")
+
 
-def load_sbom(path: Path) -> tuple[str, bytes]:
-    """Read the SBOM; `path` may be the file or a directory containing one."""
+def scan_dir(path: Path) -> Path:
+    """If `path` is a directory, return the path of the SBOM within."""
     if path.is_dir():
         candidates = list(path.glob("*.cdx.json"))
         if len(candidates) != 1:
             msg = f"expected exactly one *.cdx.json in {path}, found {len(candidates)}"
             raise SystemExit(msg)
-        path = candidates[0]
-    return path.name, path.read_bytes()
+        return candidates[0]
+    return path
 
 
 def main() -> None:
@@ -75,16 +79,15 @@ def main() -> None:
     )
     args = parser.parse_args()
 
-    sbom_name, sbom_bytes = load_sbom(args.sbom)
+    sbom = scan_dir(args.sbom)
 
     wheels = sorted(args.wheelhouse.glob("*.whl"))
     if not wheels:
         print(f"error: no wheels found in {args.wheelhouse}", file=sys.stderr)
         raise SystemExit(1)
 
     for wheel in wheels:
-        embed(wheel, sbom_name, sbom_bytes)
-        print(f"Embedded {sbom_name} in {wheel.name}")
+        embed(wheel, sbom)
 
 
 if __name__ == "__main__":