Pass the artifacts checksums in the headers of the deploy request

aslheyrr · aslheyrr · commit 6f553d301389 · 2024-09-24T17:02:27.000-05:00
diff --git a/README.md b/README.md
@@ -34,6 +34,7 @@ This library enables you to manage Artifactory resources such as users, groups,
     + [Deploy an artifact](#deploy-an-artifact)
     + [Deploy an artifact with properties](#deploy-an-artifact-with-properties)
     + [Deploy an artifact with checksums](#deploy-an-artifact-with-checksums)
+    + [Deploy an artifact by checksums](#deploy-an-artifact-by-checksums)
     + [Download an artifact](#download-an-artifact)
     + [Retrieve artifact list](#retrieve-artifact-list)
     + [Retrieve artifact properties](#retrieve-artifact-properties)
@@ -426,6 +427,16 @@ artifact = art.artifacts.deploy("<LOCAL_FILE_LOCATION>", "<ARTIFACT_PATH_IN_ARTI
 
 #### Deploy an artifact with checksums
 
+```python
+artifact = art.artifacts.deploy("<LOCAL_FILE_LOCATION>", "<ARTIFACT_PATH_IN_ARTIFACTORY>", "<CHECKSUM_ALGORITHMS>")
+# artifact = art.artifacts.deploy("Desktop/myNewFile.txt", "my-repository/my/new/artifact/directory/file.txt", checksum_algorithms=["md5", "sha1", "sha256"] )
+```
+
+Provides locally calculated checksums for files during deployment so that Artifactory can verify the authenticity of artifacts.
+
+
+#### Deploy an artifact by checksums
+
 ```python
 artifact = art.artifacts.deploy("<LOCAL_FILE_LOCATION>", "<ARTIFACT_PATH_IN_ARTIFACTORY>", checksum_enabled=True)
 # artifact = art.artifacts.deploy("Desktop/myNewFile.txt", "my-repository/my/new/artifact/directory/file.txt", checksums=True)
diff --git a/pyartifactory/exception.py b/pyartifactory/exception.py
@@ -58,3 +58,7 @@ class InvalidTokenDataError(ArtifactoryError):
 
 class BuildNotFoundError(ArtifactoryError):
     """Requested build were not found"""
+
+
+class InvalidAlgorithmError(ArtifactoryError):
+    """The Algoritnm is not supported by Artifactory."""
diff --git a/pyartifactory/models/artifact.py b/pyartifactory/models/artifact.py
@@ -10,28 +10,33 @@
 
 from pydantic import BaseModel
 
+from pyartifactory.exception import InvalidAlgorithmError
+
 
 class Checksums(BaseModel):
     """Models a checksum."""
 
-    sha1: str
-    md5: str
-    sha256: str
+    sha1: Optional[str] = None
+    md5: Optional[str] = None
+    sha256: Optional[str] = None
 
     @classmethod
-    def generate(cls, file_: Path) -> Checksums:
+    def generate(cls, file_: Path, algorithms: List[str] = ["sha1", "sha256", "md5"]) -> Checksums:
         block_size: int = 65536
         mapping: dict[str, Callable[[], Any]] = {"md5": hashlib.md5, "sha1": hashlib.sha1, "sha256": hashlib.sha256}
         results = {}
 
-        for algorithm, hashing_function in mapping.items():
-            hasher = hashing_function()
-            with file_.absolute().open("rb") as fd:
-                buf = fd.read(block_size)
-                while len(buf) > 0:
-                    hasher.update(buf)
+        for algorithm in algorithms:
+            if algorithm in mapping:
+                hasher = mapping[algorithm]()
+                with file_.absolute().open("rb") as fd:
                     buf = fd.read(block_size)
-            results[algorithm] = hasher.hexdigest()
+                    while len(buf) > 0:
+                        hasher.update(buf)
+                        buf = fd.read(block_size)
+                results[algorithm] = hasher.hexdigest()
+            else:
+                raise InvalidAlgorithmError(f"'{algorithm}' is not a supported checksum algorithm")
 
         return cls(**results)
 
@@ -40,6 +45,8 @@ class OriginalChecksums(BaseModel):
     """Models original checksums."""
 
     sha256: str
+    sha1: Optional[str] = None
+    md5: Optional[str] = None
 
 
 class Child(BaseModel):
diff --git a/pyartifactory/objects/artifact.py b/pyartifactory/objects/artifact.py
@@ -82,6 +82,7 @@ def deploy(
         artifact_path: Union[Path, str],
         properties: Optional[Dict[str, List[str]]] = None,
         checksum_enabled: bool = False,
+        checksum_algorithms: Optional[List[str]] = None,
     ) -> ArtifactInfoResponse:
         """
         Deploy a file or directory.
@@ -98,17 +99,30 @@ def deploy(
                 for file in files:
                     self.deploy(Path(f"{root}/{file}"), Path(f"{new_root}/{file}"), properties, checksum_enabled)
         else:
+            checksum_headers: Dict[str, str] = {}
+
+            if checksum_enabled:
+                checksum_headers["X-Checksum-Deploy"] = "true"
+                checksum_algorithms = ["sha1", "sha256", "md5"]
+
+            if checksum_algorithms is not None:
+                artifact_check_sums = Checksums.generate(local_file, checksum_algorithms)
+                if artifact_check_sums.md5:
+                    checksum_headers["X-Checksum"] = artifact_check_sums.md5
+                if artifact_check_sums.sha1:
+                    checksum_headers["X-Checksum-Sha1"] = artifact_check_sums.sha1
+                if artifact_check_sums.sha256:
+                    checksum_headers["X-Checksum-Sha256"] = artifact_check_sums.sha256
+
             if checksum_enabled:
-                artifact_check_sums = Checksums.generate(local_file)
                 try:
+                    properties_param_str = ""
+                    if properties is not None:
+                        properties_param_str = self._format_properties(properties)
+                    route = ";".join(s for s in [artifact_folder.as_posix(), properties_param_str] if s)
                     self._put(
-                        route=artifact_folder.as_posix(),
-                        headers={
-                            "X-Checksum-Deploy": "true",
-                            "X-Checksum-Sha1": artifact_check_sums.sha1,
-                            "X-Checksum-Sha256": artifact_check_sums.sha256,
-                            "X-Checksum": artifact_check_sums.md5,
-                        },
+                        route=route,
+                        headers=checksum_headers,
                     )
                 except requests.exceptions.HTTPError as error:
                     if error.response.status_code == 404:
@@ -125,7 +139,7 @@ def deploy(
                     if properties is not None:
                         properties_param_str = self._format_properties(properties)
                     route = ";".join(s for s in [artifact_folder.as_posix(), properties_param_str] if s)
-                    self._put(route, data=stream)
+                    self._put(route=route, headers=checksum_headers, data=stream)
 
             logger.debug("Artifact %s successfully deployed", local_file)
         return self.info(artifact_folder)
diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py
@@ -16,6 +16,7 @@
     ArtifactListFolderResponse,
     ArtifactListResponse,
     Checksums,
+    OriginalChecksums,
 )
 
 URL = "http://localhost:8080/artifactory"
@@ -609,6 +610,113 @@ def test_checksum_defined_file(file_path: Path, expected_sha1: str, expected_md5
     assert result == expected
 
 
+@pytest.mark.parametrize(
+    "file_info,expected_checksums",
+    [
+        pytest.param(
+            {
+                **FILE_INFO_RESPONSE.copy(),
+                "originalChecksums": {
+                    "md5": "4cf609e0fe1267df8815bc650f5851e9",
+                    "sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858",
+                },
+            },
+            {
+                "md5": "4cf609e0fe1267df8815bc650f5851e9",
+                "sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858",
+            },
+            id="md5",
+        ),
+        pytest.param(
+            {
+                **FILE_INFO_RESPONSE.copy(),
+                "originalChecksums": {
+                    "sha1": "962c287c760e03b03c17eb920f5358d05f44dd3b",
+                    "sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858",
+                },
+            },
+            {
+                "sha1": "962c287c760e03b03c17eb920f5358d05f44dd3b",
+                "sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858",
+            },
+            id="sha1",
+        ),
+        pytest.param(
+            {
+                **FILE_INFO_RESPONSE.copy(),
+                "originalChecksums": {"sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858"},
+            },
+            {"sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858"},
+            id="sha256",
+        ),
+        pytest.param(
+            {
+                **FILE_INFO_RESPONSE.copy(),
+                "originalChecksums": {
+                    "md5": "4cf609e0fe1267df8815bc650f5851e9",
+                    "sha1": "962c287c760e03b03c17eb920f5358d05f44dd3b",
+                    "sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858",
+                },
+            },
+            {
+                "sha256": "396cf16e8ce000342c95ffc7feb2a15701d0994b70c1b13fea7112f85ac8e858",
+                "md5": "4cf609e0fe1267df8815bc650f5851e9",
+                "sha1": "962c287c760e03b03c17eb920f5358d05f44dd3b",
+            },
+            id="md5&sha1",
+        ),
+    ],
+)
+@responses.activate
+def test_deploy_artifact_with_checksum_algorithms_success(file_info: dict, expected_checksums: dict):
+    responses.add(responses.PUT, f"{URL}/{ARTIFACT_PATH}", status=200)
+    responses.add(
+        responses.GET,
+        f"{URL}/api/storage/{ARTIFACT_PATH}",
+        json=file_info,
+        status=200,
+    )
+    expected = OriginalChecksums(**expected_checksums)
+    artifactory = ArtifactoryArtifact(AuthModel(url=URL, auth=AUTH))
+    artifact = artifactory.deploy(
+        Path(LOCAL_FILE_LOCATION),
+        Path(ARTIFACT_PATH),
+        checksum_algorithms=file_info["originalChecksums"].keys(),
+    )
+    assert expected == artifact.originalChecksums
+
+
+@pytest.mark.parametrize(
+    "file_info",
+    [
+        pytest.param(
+            {
+                **FILE_INFO_RESPONSE.copy(),
+                "originalChecksums": {"sha224": "d14a028c2a3a2bc9476102bb288234c415a2b01f828ea62ac5b3e42f"},
+            },
+            id="sha224",
+        ),
+    ],
+)
+@responses.activate
+def test_deploy_artifact_with_checksum_algorithms_error(file_info: dict):
+    responses.add(responses.PUT, f"{URL}/{ARTIFACT_PATH}", status=409)
+
+    responses.add(
+        responses.GET,
+        f"{URL}/api/storage/{ARTIFACT_PATH}",
+        json=file_info,
+        status=200,
+    )
+    with pytest.raises(Exception):
+        artifactory = ArtifactoryArtifact(AuthModel(url=URL, auth=AUTH))
+        artifactory.deploy(
+            Path(LOCAL_FILE_LOCATION),
+            Path(ARTIFACT_PATH),
+            checksum_algorithms=file_info["originalChecksums"].keys(),
+        )
+
+
 @responses.activate
 def test_deploy_artifact_with_checksum_success(mocker):
     responses.add(responses.PUT, f"{URL}/{ARTIFACT_PATH}", status=200)
