Added PEM read/write functions.

sk/certificates.py

@@ -1,10 +1,90 @@
 import datetime
-from cryptography.x509 import Name, NameAttribute, CertificateBuilder, BasicConstraints, random_serial_number
+from io import BufferedWriter
-from cryptography.x509.oid import NameOID
+import re
+import cryptography.x509 as x509
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
 import pathlib
 
+def read_private_key(pem_path: str) -> PrivateKeyTypes:
+    """
+    Read a PEM file and extract the private key in bytes.
+    """
+    with open(pem_path, "rb") as f:
+        pem = f.read()
+        key_pem = re.search(b"-----BEGIN (?:RSA )?PRIVATE KEY-----.*?END (?:RSA )?PRIVATE KEY-----", pem, re.S).group(0)
+        key = serialization.load_pem_private_key(key_pem, password=None)
+        return key
+
+def read_public_certificate(pem_path: str) -> any:
+    """
+    Read a PEM file and extract the public certificate.
+    """
+    with open(pem_path, "rb") as f:
+        pem = f.read()
+        cert_pem = re.search(b"-----BEGIN CERTIFICATE-----.*?END CERTIFICATE-----", pem, re.S).group(0)
+        cert = x509.load_pem_x509_certificate(cert_pem)
+        return cert
+
+def _write_pem_key(
+        writer: BufferedWriter,
+        key: PrivateKeyTypes,
+        encoding: serialization.Encoding = serialization.Encoding.PEM,
+        format: serialization.PrivateFormat = serialization.PrivateFormat.PKCS8,
+        password: str | None = None
+):
+    """
+    Write a PEM encoded private key to the given writer.
+    Allows optional password protection.
+    """
+    if password:
+        encryption_algorithm = serialization.BestAvailableEncryption(password.encode())
+    else:
+        encryption_algorithm = serialization.NoEncryption()
+
+    key_pem = key.private_bytes(
+        encoding=encoding,
+        format=format,
+        encryption_algorithm=encryption_algorithm,
+    )
+    writer.write(key_pem)
+
+def _write_pem_cert(
+        writer: BufferedWriter,
+        cert: x509.Certificate,
+        encoding: serialization.Encoding = serialization.Encoding.PEM
+        ):
+    cert_pem = cert.public_bytes(encoding=encoding)
+    writer.write(cert_pem)
+
+def write_pem_file(
+        pem_file: pathlib.Path,
+        cert: x509.Certificate | None = None,
+        key: PrivateKeyTypes | None = None,
+        encoding: serialization.Encoding = serialization.Encoding.PEM,
+        key_serialization_format: serialization.PrivateFormat = serialization.PrivateFormat.PKCS8,
+        password: str | None = None
+        ):
+    """
+    Write the certificate and/or private key to a PEM file.
+    """
+    with open(pem_file, "wb") as f:
+        if cert:
+            _write_pem_cert(
+                f,
+                cert,
+                encoding=encoding
+                )
+        if key:
+            _write_pem_key(
+                f,
+                key,
+                encoding=encoding,
+                password=password,
+                format=key_serialization_format
+                )
+
 def create_self_signed_certificate(
         file_path: str,
         subject_name: str,
@@ -39,46 +119,41 @@ def create_self_signed_certificate(
     """
     key = rsa.generate_private_key(public_exponent=65537, key_size=key_size)
 
-    subject = issuer = Name([
+    subject = issuer = x509.Name([
-        NameAttribute(NameOID.COUNTRY_NAME, country_name),
+        x509.oid.NameAttribute(x509.oid.NameOID.COUNTRY_NAME, country_name),
-        NameAttribute(NameOID.ORGANIZATION_NAME, organization_name),
+        x509.oid.NameAttribute(x509.oid.NameOID.ORGANIZATION_NAME, organization_name),
-        NameAttribute(NameOID.COMMON_NAME, subject_name),
+        x509.oid.NameAttribute(x509.oid.NameOID.COMMON_NAME, subject_name),
     ])
 
     cert = (
-        CertificateBuilder()
+        x509.CertificateBuilder()
         .subject_name(subject)
         .issuer_name(issuer)
         .public_key(key.public_key())
-        .serial_number(random_serial_number())
+        .serial_number(x509.random_serial_number())
         .not_valid_before(datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=1))
         .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=valid_days))
-        .add_extension(BasicConstraints(ca=True, path_length=None), critical=True)
+        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
         .sign(key, hashes.SHA256())
     )
 
     crt_path = pathlib.Path(file_path).with_suffix('.crt')
     key_path = pathlib.Path(file_path).with_suffix('.key')
     pem_path = pathlib.Path(file_path).with_suffix('.pem')
-    
+
     public_key = cert.public_bytes(serialization.Encoding.PEM)
     private_key = key.private_bytes(
         encoding=serialization.Encoding.PEM,
         format=serialization.PrivateFormat.TraditionalOpenSSL,
         encryption_algorithm=serialization.NoEncryption(),
     )
-    
-    # Write the certificate
-    with open(crt_path, "wb") as f:
-        f.write(public_key)
 
+    # Write the certificate
+    write_pem_file(crt_path, cert=cert)
     # Write the private key
-    with open(key_path, "wb") as f:
+    write_pem_file(key_path, key=key)
-        f.write(private_key)
+    # Write both to a combined PEM file
-    
+    write_pem_file(pem_path, cert=cert, key=key)
-    with open(pem_path, "wb") as f:
-        f.write(public_key)
-        f.write(private_key)
 
     print(f"Certificate created and saved.")
     print(f" - Certificate: {crt_path}")