From 367e1fffc395de91f5e3130974a5406420cdf607 Mon Sep 17 00:00:00 2001
From: Jaideep Rao <jrao@redhat.com>
Date: Wed, 8 Jan 2025 02:00:51 +0530
Subject: [PATCH] add index and oci-layout file

Signed-off-by: Jaideep Rao <jrao@redhat.com>
---
 olot/oci/oci_image_index.py  | 23 ++++++++++++++++++++-
 olot/oci/oci_image_layout.py |  8 +++++++-
 olot/oci_artifact.py         | 39 +++++++++++++++++++++++++++++++++---
 3 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/olot/oci/oci_image_index.py b/olot/oci/oci_image_index.py
index e3c76e8..2f72c9f 100644
--- a/olot/oci/oci_image_index.py
+++ b/olot/oci/oci_image_index.py
@@ -9,7 +9,7 @@
 
 from pydantic import BaseModel, Field
 
-from olot.oci.oci_common import MediaType, Digest, Urls
+from olot.oci.oci_common import MediaTypes, MediaType, Digest, Urls
 from olot.utils.types import  Int64, Base64, Annotations
 
 class Platform(BaseModel):
@@ -199,3 +199,24 @@ def read_ocilayout_root_index(ocilayout: Path) -> OCIImageIndex:
     with open(ocilayout / "index.json", "r") as f:
         ocilayout_root_index = OCIImageIndex.model_validate_json(f.read())
     return ocilayout_root_index
+
+
+def create_oci_image_index(
+    schemaVersion: int = 2,
+    mediaType: Optional[str] = MediaTypes.index,
+    artifactType: Optional[str] = None,
+    subject: Optional[ContentDescriptor] = None,
+    manifests: List[Manifest] = [],
+    annotations: Optional[Annotations] = None
+) -> OCIImageIndex:
+    """
+    Create an OCI image index object.
+    """
+    return OCIImageIndex(
+        schemaVersion=schemaVersion,
+        mediaType=mediaType,
+        artifactType=artifactType,
+        subject=subject,
+        manifests=manifests,
+        annotations=annotations
+    )
\ No newline at end of file
diff --git a/olot/oci/oci_image_layout.py b/olot/oci/oci_image_layout.py
index 799db7d..6c50326 100644
--- a/olot/oci/oci_image_layout.py
+++ b/olot/oci/oci_image_layout.py
@@ -19,10 +19,16 @@ class OCIImageLayout(BaseModel):
         ..., description='version of the OCI Image Layout (in the oci-layout file)'
     )
 
+    class Config:
+        use_enum_values = True
+
 def verify_ocilayout(ocilayout: Path):
     with open(ocilayout / "oci-layout", "r") as f:
         m = OCIImageLayout.model_validate_json(f.read())
-        if not m.imageLayoutVersion == ImageLayoutVersion.field_1_0_0:
+        if not m.imageLayoutVersion == ImageLayoutVersion.field_1_0_0.value:
             raise ValueError(f"Unexpected ocilayout in {ocilayout}")
         else:
             return True
+
+def create_ocilayout() -> OCIImageLayout:
+    return OCIImageLayout(imageLayoutVersion=ImageLayoutVersion.field_1_0_0)
\ No newline at end of file
diff --git a/olot/oci_artifact.py b/olot/oci_artifact.py
index 9150679..709dded 100644
--- a/olot/oci_artifact.py
+++ b/olot/oci_artifact.py
@@ -6,11 +6,20 @@
 from typing import List
 
 from olot.oci.oci_image_manifest import create_oci_image_manifest, create_manifest_layers
-from olot.oci.oci_common import Keys
+from olot.oci.oci_image_layout import create_ocilayout
+from olot.oci.oci_common import Keys, MediaTypes
+from olot.oci.oci_image_index import Manifest, create_oci_image_index
 from olot.utils.files import MIMETypes, tarball_from_file, targz_from_file
 from olot.utils.types import compute_hash_of_str
 
 def create_oci_artifact_from_model(source_dir: Path, dest_dir: Path):
+    """
+    Create an OCI artifact from a model directory.
+
+    Args:
+        source_dir: The directory containing the model files.
+        dest_dir: The directory to write the OCI artifact to. If None, a directory named 'oci' will be created in the source directory.
+    """
     if not source_dir.exists():
         raise NotADirectoryError(f"Input directory '{source_dir}' does not exist.")
 
@@ -39,11 +48,34 @@ def create_oci_artifact_from_model(source_dir: Path, dest_dir: Path):
         layers=manifest_layers,
         annotations=annotations
     )
-    manifest_json = json.dumps(manifest.dict(), indent=4, sort_keys=True)
+    manifest_json = json.dumps(manifest.model_dump(), indent=4, sort_keys=True)
     manifest_SHA = compute_hash_of_str(manifest_json)
     with open(sha256_path / manifest_SHA, "w") as f:
         f.write(manifest_json)
 
+    # Create the OCI image index
+    index = create_oci_image_index(
+        manifests = [
+            Manifest(
+                mediaType=MediaTypes.manifest,
+                size=os.stat(sha256_path / manifest_SHA).st_size,
+                digest=f"sha256:{manifest_SHA}",
+                urls=None,
+                platform=None,
+                annotations=None
+            )
+        ]
+    )
+    index_json = json.dumps(index.model_dump(), indent=4, sort_keys=True)
+    with open(dest_dir / "index.json", "w") as f:
+        f.write(index_json)
+
+
+    # Create the OCI-layout file
+    oci_layout = create_ocilayout()
+    with open(dest_dir / "oci-layout", "w") as f:
+        f.write(json.dumps(oci_layout.model_dump(), indent=4, sort_keys=True))
+
 
 def create_blobs(model_files: List[Path], dest_dir: Path):
     """
@@ -72,4 +104,5 @@ def main():
     source_dir = Path(args.source_dir)
     create_oci_artifact_from_model(source_dir, None)
 
-
+if __name__ == "__main__":
+    main()