Allow multiple certs in bundle

someengineering · Oct 2, 2023 · e73c41e · e73c41e
1 parent b583688
commit e73c41e
Showing 1 changed file with 21 additions and 16 deletions.
diff --git a/resotolib/resotolib/x509.py b/resotolib/resotolib/x509.py
@@ -186,30 +186,35 @@ def write_cert_to_file(cert: Certificate, cert_path: str, rename: bool = True) -
         os.rename(tmp_cert_path, cert_path)
 
 
-def gen_ca_bundle_bytes(cert: Certificate, include_certifi: bool = True) -> bytes:
+def gen_ca_bundle_bytes(certs: Union[Certificate, List[Certificate]], include_certifi: bool = True) -> bytes:
     content = bytearray()
     if include_certifi:
         content.extend(certifi.contents().encode())
-    content.extend("\n".encode())
-    content.extend(f"# Issuer: {cert.issuer.rfc4514_string()}\n".encode())
-    content.extend(f"# Subject: {cert.subject.rfc4514_string()}\n".encode())
-    label: str = cert.issuer.get_attributes_for_oid(NameOID.COMMON_NAME)[0].value  # type: ignore
-    content.extend(f"# Label: {label}\n".encode())
-    content.extend(f"# Serial: {cert.serial_number}\n".encode())
-    md5 = cert_fingerprint(cert, "MD5")
-    sha1 = cert_fingerprint(cert, "SHA1")
-    sha256 = cert_fingerprint(cert, "SHA256")
-    content.extend(f"# MD5 Fingerprint: {md5}\n".encode())
-    content.extend(f"# SHA1 Fingerprint: {sha1}\n".encode())
-    content.extend(f"# SHA256 Fingerprint: {sha256}\n".encode())
-    content.extend(cert_to_bytes(cert))
+
+    if isinstance(certs, Certificate):
+        certs = [certs]
+
+    for cert in certs:
+        content.extend("\n".encode())
+        content.extend(f"# Issuer: {cert.issuer.rfc4514_string()}\n".encode())
+        content.extend(f"# Subject: {cert.subject.rfc4514_string()}\n".encode())
+        label: str = cert.issuer.get_attributes_for_oid(NameOID.COMMON_NAME)[0].value  # type: ignore
+        content.extend(f"# Label: {label}\n".encode())
+        content.extend(f"# Serial: {cert.serial_number}\n".encode())
+        md5 = cert_fingerprint(cert, "MD5")
+        sha1 = cert_fingerprint(cert, "SHA1")
+        sha256 = cert_fingerprint(cert, "SHA256")
+        content.extend(f"# MD5 Fingerprint: {md5}\n".encode())
+        content.extend(f"# SHA1 Fingerprint: {sha1}\n".encode())
+        content.extend(f"# SHA256 Fingerprint: {sha256}\n".encode())
+        content.extend(cert_to_bytes(cert))
     return bytes(content)
 
 
-def write_ca_bundle(cert: Certificate, cert_path: str, include_certifi: bool = True, rename: bool = True) -> None:
+def write_ca_bundle(certs: Union[Certificate, List[Certificate]], cert_path: str, include_certifi: bool = True, rename: bool = True) -> None:
     tmp_cert_path = f"{cert_path}.tmp" if rename else cert_path
     with open(tmp_cert_path, "wb") as f:
-        f.write(gen_ca_bundle_bytes(cert, include_certifi))
+        f.write(gen_ca_bundle_bytes(certs, include_certifi))
     if rename:
         os.rename(tmp_cert_path, cert_path)