From 81f8cf59b1f40ffbd213789a8b1b621a01e3f631 Mon Sep 17 00:00:00 2001
From: Jan Kowalleck <jan.kowalleck@gmail.com>
Date: Fri, 22 Dec 2023 02:43:00 +0100
Subject: [PATCH] feat: add function to map python `hashlib` algorithms to
 CycloneDX (#519)

new API: `model.HashType.from_hashlib_alg()`

Signed-off-by: Jan Kowalleck <jan.kowalleck@gmail.com>
---
 cyclonedx/model/__init__.py | 43 +++++++++++++++++++++++++++++++++++++
 tests/test_model.py         | 41 ++++++++++++++++++++++++-----------
 2 files changed, 72 insertions(+), 12 deletions(-)

diff --git a/cyclonedx/model/__init__.py b/cyclonedx/model/__init__.py
index 95f26ccb..5c7726a4 100644
--- a/cyclonedx/model/__init__.py
+++ b/cyclonedx/model/__init__.py
@@ -343,6 +343,25 @@ def xml_denormalize(cls, o: 'XmlElement', *,
         ]
 
 
+_MAP_HASHLIB: Dict[str, HashAlgorithm] = {
+    # from hashlib.algorithms_guaranteed
+    'md5': HashAlgorithm.MD5,
+    'sha1': HashAlgorithm.SHA_1,
+    # sha224:
+    'sha256': HashAlgorithm.SHA_256,
+    'sha384': HashAlgorithm.SHA_384,
+    'sha512': HashAlgorithm.SHA_512,
+    # blake2b:
+    # blake2s:
+    # sha3_224:
+    'sha3_256': HashAlgorithm.SHA3_256,
+    'sha3_384': HashAlgorithm.SHA3_384,
+    'sha3_512': HashAlgorithm.SHA3_512,
+    # shake_128:
+    # shake_256:
+}
+
+
 @serializable.serializable_class
 class HashType:
     """
@@ -352,6 +371,30 @@ class HashType:
         See the CycloneDX Schema for hashType: https://cyclonedx.org/docs/1.3/#type_hashType
     """
 
+    @staticmethod
+    def from_hashlib_alg(hashlib_alg: str, content: str) -> 'HashType':
+        """
+        Attempts to convert a hashlib-algorithm to our internal model classes.
+
+        Args:
+             hashlib_alg:
+                Hash algorith - like it is used by `hashlib`.
+                Example: `sha256`.
+
+            content:
+                Hash value.
+
+        Raises:
+            `UnknownHashTypeException` if the algorithm of hash cannot be determined.
+
+        Returns:
+            An instance of `HashType`.
+        """
+        alg = _MAP_HASHLIB.get(hashlib_alg.lower())
+        if alg is None:
+            raise UnknownHashTypeException(f'Unable to determine hash alg for {hashlib_alg!r}')
+        return HashType(alg=alg, content=content)
+
     @staticmethod
     def from_composite_str(composite_hash: str) -> 'HashType':
         """
diff --git a/tests/test_model.py b/tests/test_model.py
index dc8a6ff0..e815bdc2 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -20,6 +20,8 @@
 from enum import Enum
 from unittest import TestCase
 
+from ddt import ddt, named_data
+
 from cyclonedx._internal.compare import ComparableTuple
 from cyclonedx.exception.model import (
     InvalidLocaleTypeException,
@@ -139,8 +141,8 @@ def test_compare_last_item_missing(self) -> None:
         self.assertNotEqual(tuple2, tuple1)
 
     def test_compare_enum(self) -> None:
-        tuple1 = ComparableTuple((DummyStringEnum.FIRST, ))
-        tuple2 = ComparableTuple((DummyStringEnum.SECOND, ))
+        tuple1 = ComparableTuple((DummyStringEnum.FIRST,))
+        tuple2 = ComparableTuple((DummyStringEnum.SECOND,))
         self.assertLess(tuple1, tuple2)
         self.assertGreater(tuple2, tuple1)
         self.assertNotEqual(tuple1, tuple2)
@@ -239,19 +241,34 @@ def test_sort(self) -> None:
         self.assertListEqual(sorted_refs, expected_refs)
 
 
+@ddt
 class TestModelHashType(TestCase):
 
-    def test_hash_type_from_composite_str_1(self) -> None:
-        h = HashType.from_composite_str('sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b')
-        self.assertEqual(h.alg, HashAlgorithm.SHA_256)
-        self.assertEqual(h.content, '806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b')
-
-    def test_hash_type_from_composite_str_2(self) -> None:
-        h = HashType.from_composite_str('md5:dc26cd71b80d6757139f38156a43c545')
-        self.assertEqual(h.alg, HashAlgorithm.MD5)
-        self.assertEqual(h.content, 'dc26cd71b80d6757139f38156a43c545')
+    @named_data(
+        ('sha256', 'sha256', '806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b', HashAlgorithm.SHA_256),
+        ('MD5', 'MD5', 'dc26cd71b80d6757139f38156a43c545', HashAlgorithm.MD5),
+    )
+    def test_hash_type_from_hashlib_alg(self, alg: str, content: str, e_alg: HashAlgorithm) -> None:
+        h = HashType.from_hashlib_alg(alg, content)
+        self.assertIs(h.alg, e_alg)
+        self.assertEqual(h.content, content)
 
-    def test_hash_type_from_unknown(self) -> None:
+    def test_hash_type_from_hashlib_alg_throws_on_unknown(self) -> None:
+        with self.assertRaises(UnknownHashTypeException):
+            HashType.from_hashlib_alg('unknown', 'dc26cd71b80d6757139f38156a43c545')
+
+    @named_data(
+        ('sha256', 'sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b',
+         HashAlgorithm.SHA_256, '806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b'),
+        ('MD5', 'MD5:dc26cd71b80d6757139f38156a43c545',
+         HashAlgorithm.MD5, 'dc26cd71b80d6757139f38156a43c545'),
+    )
+    def test_hash_type_from_composite_str(self, composite: str, e_alg: HashAlgorithm, e_content: str) -> None:
+        h = HashType.from_composite_str(composite)
+        self.assertIs(h.alg, e_alg)
+        self.assertEqual(h.content, e_content)
+
+    def test_hash_type_from_composite_str_throws_on_unknown(self) -> None:
         with self.assertRaises(UnknownHashTypeException):
             HashType.from_composite_str('unknown:dc26cd71b80d6757139f38156a43c545')