diff --git a/oras/defaults.py b/oras/defaults.py index a1452e24..b2a98449 100644 --- a/oras/defaults.py +++ b/oras/defaults.py @@ -48,3 +48,6 @@ class registry: blank_config_hash = ( "sha256:44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a" ) + +# ModelPack annotations, refer to https://github.com/modelpack/model-spec/blob/main/specs-go/v1/annotations.go +annotation_filepath = "org.cnai.model.filepath" diff --git a/oras/provider.py b/oras/provider.py index cb0cfc9b..e8d8f318 100644 --- a/oras/provider.py +++ b/oras/provider.py @@ -802,6 +802,12 @@ def push( if annotations: layer["annotations"].update(annotations) + # if filepath annotation is present, set it + if oras.defaults.annotation_filepath in layer["annotations"]: + layer["annotations"][oras.defaults.annotation_filepath] = ( + blob_name.strip(os.sep) + ) + # update the manifest with the new layer manifest["layers"].append(layer) logger.debug(f"Preparing layer {layer}") @@ -901,8 +907,8 @@ def pull( files = [] for layer in manifest.get("layers", []): filename = (layer.get("annotations") or {}).get( - oras.defaults.annotation_title - ) + oras.defaults.annotation_filepath + ) or (layer.get("annotations") or {}).get(oras.defaults.annotation_title) # If we don't have a filename, default to digest. Hopefully does not happen if not filename: @@ -917,18 +923,32 @@ def pull( ) continue - # A directory will need to be uncompressed and moved - if layer["mediaType"] == oras.defaults.default_blob_dir_media_type: - targz = oras.utils.get_tmpfile(suffix=".tar.gz") - self.download_blob(container, layer["digest"], targz) - - # The artifact will be extracted to the correct name - oras.utils.extract_targz(targz, os.path.dirname(outfile)) + # Determine compression format from mediaType and handle accordingly + media_type = layer["mediaType"] + compression = oras.utils.get_compression_from_media_type(media_type) + + # Handle compressed archives that need extraction + if compression in ["gzip", "zstd", "tar"]: + # Get appropriate file extension based on compression + if compression == "gzip": + suffix = ".tar.gz" + elif compression == "zstd": + suffix = ".tar.zst" + else: # tar + suffix = ".tar" + + archive_file = oras.utils.get_tmpfile(suffix=suffix) + self.download_blob(container, layer["digest"], archive_file) + + # Extract the archive to the correct location + oras.utils.extract_by_compression( + archive_file, os.path.dirname(outfile), compression + ) # Anything else just extracted directly else: self.download_blob(container, layer["digest"], outfile) - logger.info(f"Successfully pulled {outfile}.") + logger.info(f"Successfully pulled {outfile}") files.append(outfile) return files diff --git a/oras/utils/__init__.py b/oras/utils/__init__.py index a64749c7..acb78c3e 100644 --- a/oras/utils/__init__.py +++ b/oras/utils/__init__.py @@ -1,6 +1,10 @@ from .fileio import ( copyfile, + extract_by_compression, + extract_tar, + extract_tar_zstd, extract_targz, + get_compression_from_media_type, get_file_hash, get_size, get_tmpdir, diff --git a/oras/utils/fileio.py b/oras/utils/fileio.py index 294f11e1..f96e6535 100644 --- a/oras/utils/fileio.py +++ b/oras/utils/fileio.py @@ -103,6 +103,92 @@ def extract_targz(targz: str, outdir: str, numeric_owner: bool = False): tar.extractall(outdir, members=None, numeric_owner=numeric_owner) +def extract_tar(tar_file: str, outdir: str, numeric_owner: bool = False): + """ + Extract a .tar (uncompressed) to an output directory. + """ + with tarfile.open(tar_file, "r:") as tar: + for member in tar.getmembers(): + member_path = os.path.join(outdir, member.name) + if not is_within_directory(outdir, member_path): + raise Exception("Attempted Path Traversal in Tar File") + tar.extractall(outdir, members=None, numeric_owner=numeric_owner) + + +def extract_tar_zstd(tar_zstd: str, outdir: str, numeric_owner: bool = False): + """ + Extract a .tar.zst/.tar.zstd to an output directory. + Requires zstandard package to be installed. + """ + try: + import zstandard as zstd + except ImportError: + raise ImportError( + "zstandard package is required for zstd decompression. " + "Install it with: pip install zstandard" + ) + + with open(tar_zstd, "rb") as compressed_file: + dctx = zstd.ZstdDecompressor() + with dctx.stream_reader(compressed_file) as reader: + with tarfile.open(fileobj=reader, mode="r|") as tar: + for member in tar: + member_path = os.path.join(outdir, member.name) + if not is_within_directory(outdir, member_path): + raise Exception("Attempted Path Traversal in Tar File") + tar.extract(member, outdir, numeric_owner=numeric_owner) + + +def get_compression_from_media_type(media_type: str) -> str: + """ + Determine compression format from media type suffix. + + :param media_type: The media type string + :type media_type: str + :return: Compression format ('gzip', 'zstd', 'tar', or 'raw') + :rtype: str + """ + if media_type.endswith("+gzip"): + return "gzip" + elif media_type.endswith("+zstd"): + return "zstd" + elif media_type.endswith(".tar"): + return "tar" + elif media_type.endswith(".raw"): + return "raw" + else: + # Default to gzip for backward compatibility + return "gzip" + + +def extract_by_compression( + archive_file: str, outdir: str, compression: str, numeric_owner: bool = False +): + """ + Extract archive based on compression format. + + :param archive_file: Path to the archive file + :type archive_file: str + :param outdir: Output directory for extraction + :type outdir: str + :param compression: Compression format ('gzip', 'zstd', 'tar', or 'raw') + :type compression: str + :param numeric_owner: Whether to use numeric owner + :type numeric_owner: bool + """ + if compression == "gzip": + extract_targz(archive_file, outdir, numeric_owner) + elif compression == "zstd": + extract_tar_zstd(archive_file, outdir, numeric_owner) + elif compression == "tar": + extract_tar(archive_file, outdir, numeric_owner) + elif compression == "raw": + # For raw files, no extraction needed - they should be handled differently + raise ValueError("Raw files should not be extracted as archives") + else: + raise ValueError(f"Unsupported compression format: {compression}") + + def is_within_directory(directory: str, target: str) -> bool: """ Determine whether a file is within a directory