lint and test fixes

Author: laurafitzgerald

src/codeflare_sdk/common/utils/generate_cert.py

@@ -170,7 +170,7 @@ def generate_tls_cert(cluster_name, namespace, days=30, force_regenerate=False):
     - AuthorityKeyIdentifier
 
     Files are created with restricted permissions (0600) for security.
-    
+
     Certificates are stored in a user-private directory:
     - Default: ~/.local/share/codeflare/tls/{cluster_name}-{namespace}/
     - Override via CODEFLARE_TLS_DIR environment variable
@@ -194,27 +194,27 @@ def generate_tls_cert(cluster_name, namespace, days=30, force_regenerate=False):
     Raises:
         Exception:
             If an error occurs while retrieving the CA secret.
-    
+
     Example:
         # Normal generation
         generate_tls_cert("my-cluster", "default")
-        
+
         # Force regeneration if CA was rotated
         generate_tls_cert("my-cluster", "default", force_regenerate=True)
     """
     tls_base_dir = _get_tls_base_dir()
     tls_dir = tls_base_dir / f"{cluster_name}-{namespace}"
-    
+
     # Check if certificates already exist and skip if not forcing regeneration
     if not force_regenerate and tls_dir.exists():
         ca_crt = tls_dir / "ca.crt"
         tls_crt = tls_dir / "tls.crt"
         tls_key = tls_dir / "tls.key"
-        
+
         if ca_crt.exists() and tls_crt.exists() and tls_key.exists():
             # Certificates already exist, no need to regenerate
             return
-    
+
     # Create directory with secure permissions (including parent directories)
     tls_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
     tls_dir = str(tls_dir)
@@ -254,7 +254,6 @@ def generate_tls_cert(cluster_name, namespace, days=30, force_regenerate=False):
         f.write(base64.b64decode(ca_cert).decode("utf-8"))
     os.chmod(ca_crt_path, stat.S_IRUSR | stat.S_IWUSR)  # Set permissions to 0600
 
-
     # Generate tls.key and signed tls.cert locally for ray client
     # Similar to running these commands:
     # openssl req -nodes -newkey rsa:3072 -keyout ${TLSDIR}/tls.key -out ${TLSDIR}/tls.csr -subj '/CN=local'
@@ -325,21 +324,25 @@ def generate_tls_cert(cluster_name, namespace, days=30, force_regenerate=False):
             critical=True,
         )
         .add_extension(
-            x509.ExtendedKeyUsage([
-                ExtendedKeyUsageOID.SERVER_AUTH,
-                ExtendedKeyUsageOID.CLIENT_AUTH,  # For mTLS support
-            ]),
+            x509.ExtendedKeyUsage(
+                [
+                    ExtendedKeyUsageOID.SERVER_AUTH,
+                    ExtendedKeyUsageOID.CLIENT_AUTH,  # For mTLS support
+                ]
+            ),
             critical=True,
         )
         .add_extension(
-            x509.SubjectAlternativeName([
-                x509.DNSName("localhost"),
-                x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
-                x509.IPAddress(ipaddress.IPv6Address("::1")),
-                x509.DNSName(head_svc_name),
-                x509.DNSName(service_dns),
-                x509.DNSName(service_dns_cluster_local),
-            ]),
+            x509.SubjectAlternativeName(
+                [
+                    x509.DNSName("localhost"),
+                    x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
+                    x509.IPAddress(ipaddress.IPv6Address("::1")),
+                    x509.DNSName(head_svc_name),
+                    x509.DNSName(service_dns),
+                    x509.DNSName(service_dns_cluster_local),
+                ]
+            ),
             critical=False,
         )
         .add_extension(
@@ -362,7 +365,7 @@ def generate_tls_cert(cluster_name, namespace, days=30, force_regenerate=False):
     with open(tls_crt_path, "w") as f:
         f.write(tls_cert.public_bytes(serialization.Encoding.PEM).decode("utf-8"))
     os.chmod(tls_crt_path, stat.S_IRUSR | stat.S_IWUSR)  # Set permissions to 0600
- 
+
     del ca_key, ca_private_key
     try:
         del secret
@@ -467,27 +470,31 @@ def list_tls_certificates():
             created = datetime.datetime.fromtimestamp(stat_info.st_ctime)
 
             # Calculate total size
-            total_size = sum(f.stat().st_size for f in cert_dir.rglob('*') if f.is_file())
+            total_size = sum(
+                f.stat().st_size for f in cert_dir.rglob("*") if f.is_file()
+            )
 
             # Try to read certificate expiry
             cert_expiry = None
             tls_cert_path = cert_dir / "tls.crt"
             if tls_cert_path.exists():
                 try:
-                    with open(tls_cert_path, 'rb') as f:
+                    with open(tls_cert_path, "rb") as f:
                         cert = x509.load_pem_x509_certificate(f.read())
                         cert_expiry = cert.not_valid_after_utc
                 except Exception:
                     cert_expiry = None
 
-            certificates.append({
-                'cluster_name': cluster_name,
-                'namespace': namespace,
-                'path': str(cert_dir),
-                'created': created,
-                'size': total_size,
-                'cert_expiry': cert_expiry,
-            })
+            certificates.append(
+                {
+                    "cluster_name": cluster_name,
+                    "namespace": namespace,
+                    "path": str(cert_dir),
+                    "created": created,
+                    "size": total_size,
+                    "cert_expiry": cert_expiry,
+                }
+            )
 
     return certificates
 
@@ -508,7 +515,7 @@ def cleanup_expired_certificates(dry_run=True):
         >>> # Check what would be deleted
         >>> expired = cleanup_expired_certificates(dry_run=True)
         >>> print(f"Found {len(expired)} expired certificates")
-        >>> 
+        >>>
         >>> # Actually delete them
         >>> cleanup_expired_certificates(dry_run=False)
     """
@@ -520,11 +527,11 @@ def cleanup_expired_certificates(dry_run=True):
     certificates = list_tls_certificates()
 
     for cert_info in certificates:
-        if cert_info['cert_expiry'] and cert_info['cert_expiry'] < now:
-            expired_certs.append(cert_info['path'])
+        if cert_info["cert_expiry"] and cert_info["cert_expiry"] < now:
+            expired_certs.append(cert_info["path"])
 
             if not dry_run:
-                cert_dir = Path(cert_info['path'])
+                cert_dir = Path(cert_info["path"])
                 if cert_dir.exists():
                     shutil.rmtree(cert_dir)
 
@@ -534,62 +541,62 @@ def cleanup_expired_certificates(dry_run=True):
 def cleanup_old_certificates(days=30, dry_run=True):
     """
     Removes TLS certificates older than a specified number of days.
-    
+
     Args:
         days (int):
             Remove certificates created more than this many days ago. Default is 30.
         dry_run (bool):
             If True (default), only lists old certificates without deleting them.
             Set to False to actually delete old certificates.
-    
+
     Returns:
         list: List of certificate paths that were (or would be) removed.
-    
+
     Example:
         >>> # Check certificates older than 90 days
         >>> old = cleanup_old_certificates(days=90, dry_run=True)
         >>> print(f"Found {len(old)} certificates older than 90 days")
-        >>> 
+        >>>
         >>> # Delete certificates older than 30 days
         >>> cleanup_old_certificates(days=30, dry_run=False)
     """
     import shutil
-    
+
     cutoff_date = datetime.datetime.now() - datetime.timedelta(days=days)
     old_certs = []
-    
+
     certificates = list_tls_certificates()
-    
+
     for cert_info in certificates:
-        if cert_info['created'] < cutoff_date:
-            old_certs.append(cert_info['path'])
-            
+        if cert_info["created"] < cutoff_date:
+            old_certs.append(cert_info["path"])
+
             if not dry_run:
-                cert_dir = Path(cert_info['path'])
+                cert_dir = Path(cert_info["path"])
                 if cert_dir.exists():
                     shutil.rmtree(cert_dir)
-    
+
     return old_certs
 
 
 def refresh_tls_cert(cluster_name, namespace, days=30):
     """
     Refreshes TLS certificates by removing old ones and generating new ones.
-    
+
     This is useful when the server CA secret has been rotated and existing
     client certificates are no longer valid.
-    
+
     Args:
         cluster_name (str):
             The name of the Ray cluster.
         namespace (str):
             The Kubernetes namespace where the Ray cluster is located.
         days (int):
             The number of days for which the new TLS certificate will be valid. Default is 30.
-    
+
     Returns:
         bool: True if certificates were successfully refreshed.
-    
+
     Example:
         >>> # Server CA was rotated, refresh client certificates
         >>> refresh_tls_cert("my-cluster", "default")
@@ -598,12 +605,13 @@ def refresh_tls_cert(cluster_name, namespace, days=30):
     """
     # Remove old certificates
     cleanup_tls_cert(cluster_name, namespace)
-    
+
     # Generate new ones
     generate_tls_cert(cluster_name, namespace, days=days, force_regenerate=True)
-    
+
     return True
 
+
 def _get_tls_base_dir():
     """
     Get the base directory for TLS certificate storage.
@@ -617,14 +625,14 @@ def _get_tls_base_dir():
         Path: Base directory for TLS certificates
     """
     # Check for explicit override
-    tls_dir_env = os.environ.get('CODEFLARE_TLS_DIR')
+    tls_dir_env = os.environ.get("CODEFLARE_TLS_DIR")
     if tls_dir_env:
         return Path(tls_dir_env)
 
     # Use XDG Base Directory specification
-    xdg_data_home = os.environ.get('XDG_DATA_HOME')
+    xdg_data_home = os.environ.get("XDG_DATA_HOME")
     if xdg_data_home:
         return Path(xdg_data_home) / "codeflare" / "tls"
 
     # Fallback to standard location
-    return Path.home() / ".local" / "share" / "codeflare" / "tls"
+    return Path.home() / ".local" / "share" / "codeflare" / "tls"