mcp-proxy-adapter 6.3.4__py3-none-any.whl → 6.3.6__py3-none-any.whl

mcp_proxy_adapter/examples/scripts/generate_certificates_and_tokens.py

@@ -13,15 +13,24 @@ import argparse
 import subprocess
 from pathlib import Path
 from typing import Dict, Any, List, Optional
+
 # Import mcp_security_framework
 try:
     from mcp_security_framework.core.cert_manager import CertificateManager
-    from mcp_security_framework.schemas.config import CertificateConfig, CAConfig, ServerCertConfig, ClientCertConfig
+    from mcp_security_framework.schemas.config import (
+        CertificateConfig,
+        CAConfig,
+        ServerCertConfig,
+        ClientCertConfig,
+    )
     from mcp_security_framework.schemas.models import CertificateType
+
     SECURITY_FRAMEWORK_AVAILABLE = True
 except ImportError:
     SECURITY_FRAMEWORK_AVAILABLE = False
     print("Warning: mcp_security_framework not available, falling back to OpenSSL")
+
+
 def generate_ca_certificate(output_dir: str) -> str:
     """
     Generate CA certificate and key using mcp_security_framework.
@@ -44,7 +53,7 @@ def generate_ca_certificate(output_dir: str) -> str:
                 locality="City",
                 validity_years=10,  # Используем validity_years вместо validity_days
                 key_size=2048,
-                hash_algorithm="sha256"
+                hash_algorithm="sha256",
             )
             # Create certificate manager
             cert_config = CertificateConfig(
@@ -52,7 +61,7 @@ def generate_ca_certificate(output_dir: str) -> str:
                 key_storage_path=ca_dir,
                 default_validity_days=3650,
                 key_size=2048,
-                hash_algorithm="sha256"
+                hash_algorithm="sha256",
             )
             cert_manager = CertificateManager(cert_config)
             # Create CA certificate
@@ -71,16 +80,34 @@ def generate_ca_certificate(output_dir: str) -> str:
         ca_key = os.path.join(ca_dir, "ca.key")
         ca_cert = os.path.join(ca_dir, "ca.crt")
         # Generate CA private key
-        subprocess.run([
-            "openssl", "genrsa", "-out", ca_key, "2048"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            ["openssl", "genrsa", "-out", ca_key, "2048"],
+            check=True,
+            capture_output=True,
+        )
         # Generate CA certificate
-        subprocess.run([
-            "openssl", "req", "-new", "-x509", "-days", "365", "-key", ca_key,
-            "-out", ca_cert, "-subj", "/C=US/ST=State/L=City/O=Organization/CN=CA"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            [
+                "openssl",
+                "req",
+                "-new",
+                "-x509",
+                "-days",
+                "365",
+                "-key",
+                ca_key,
+                "-out",
+                ca_cert,
+                "-subj",
+                "/C=US/ST=State/L=City/O=Organization/CN=CA",
+            ],
+            check=True,
+            capture_output=True,
+        )
         print(f"✅ Generated CA certificate using OpenSSL: {ca_cert}")
         return ca_cert
+
+
 def generate_server_certificate(output_dir: str, ca_cert: str) -> tuple[str, str]:
     """
     Generate server certificate and key using mcp_security_framework.
@@ -98,10 +125,10 @@ def generate_server_certificate(output_dir: str, ca_cert: str) -> tuple[str, str
         try:
             # Find CA key file
             ca_key = None
-            if ca_cert.endswith('.crt'):
-                ca_key = ca_cert.replace('.crt', '.key')
-            elif ca_cert.endswith('.pem'):
-                ca_key = ca_cert.replace('.pem', '_key.pem')
+            if ca_cert.endswith(".crt"):
+                ca_key = ca_cert.replace(".crt", ".key")
+            elif ca_cert.endswith(".pem"):
+                ca_key = ca_cert.replace(".pem", "_key.pem")
             if not os.path.exists(ca_key):
                 print(f"❌ CA key file not found: {ca_key}")
                 return None, None
@@ -116,9 +143,12 @@ def generate_server_certificate(output_dir: str, ca_cert: str) -> tuple[str, str
                 validity_days=365,
                 key_size=2048,
                 hash_algorithm="sha256",
-                subject_alt_names=["localhost", "127.0.0.1"],  # Используем subject_alt_names вместо san_dns
+                subject_alt_names=[
+                    "localhost",
+                    "127.0.0.1",
+                ],  # Используем subject_alt_names вместо san_dns
                 ca_cert_path=ca_cert,
-                ca_key_path=ca_key
+                ca_key_path=ca_key,
             )
             # Create certificate manager
             cert_config = CertificateConfig(
@@ -126,7 +156,7 @@ def generate_server_certificate(output_dir: str, ca_cert: str) -> tuple[str, str
                 key_storage_path=keys_dir,
                 default_validity_days=365,
                 key_size=2048,
-                hash_algorithm="sha256"
+                hash_algorithm="sha256",
             )
             cert_manager = CertificateManager(cert_config)
             # Create server certificate
@@ -135,7 +165,9 @@ def generate_server_certificate(output_dir: str, ca_cert: str) -> tuple[str, str
                 print(f"✅ Generated server certificate using mcp_security_framework")
                 return (cert_pair.certificate_path, cert_pair.private_key_path)
             else:
-                print(f"❌ Failed to create server certificate: Invalid certificate pair")
+                print(
+                    f"❌ Failed to create server certificate: Invalid certificate pair"
+                )
                 return None, None
         except Exception as e:
             print(f"❌ Error creating server certificate with framework: {e}")
@@ -147,25 +179,61 @@ def generate_server_certificate(output_dir: str, ca_cert: str) -> tuple[str, str
         server_csr = os.path.join(certs_dir, "server.csr")
         server_cert = os.path.join(certs_dir, "server.crt")
         # Generate server private key
-        subprocess.run([
-            "openssl", "genrsa", "-out", server_key, "2048"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            ["openssl", "genrsa", "-out", server_key, "2048"],
+            check=True,
+            capture_output=True,
+        )
         # Generate server certificate signing request
-        subprocess.run([
-            "openssl", "req", "-new", "-key", server_key, "-out", server_csr,
-            "-subj", "/C=US/ST=State/L=City/O=Organization/CN=localhost"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            [
+                "openssl",
+                "req",
+                "-new",
+                "-key",
+                server_key,
+                "-out",
+                server_csr,
+                "-subj",
+                "/C=US/ST=State/L=City/O=Organization/CN=localhost",
+            ],
+            check=True,
+            capture_output=True,
+        )
         # Sign server certificate with CA
-        subprocess.run([
-            "openssl", "x509", "-req", "-in", server_csr, "-CA", ca_cert,
-            "-CAkey", ca_key, "-CAcreateserial", "-out", server_cert, "-days", "365"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            [
+                "openssl",
+                "x509",
+                "-req",
+                "-in",
+                server_csr,
+                "-CA",
+                ca_cert,
+                "-CAkey",
+                ca_key,
+                "-CAcreateserial",
+                "-out",
+                server_cert,
+                "-days",
+                "365",
+            ],
+            check=True,
+            capture_output=True,
+        )
         # Clean up CSR
         os.remove(server_csr)
         print(f"✅ Generated server certificate using OpenSSL: {server_cert}")
         return server_cert, server_key
-def generate_client_certificate(output_dir: str, ca_cert: str, client_name: str = "client",
-                               roles: List[str] = None, permissions: List[str] = None) -> tuple[str, str]:
+
+
+def generate_client_certificate(
+    output_dir: str,
+    ca_cert: str,
+    client_name: str = "client",
+    roles: List[str] = None,
+    permissions: List[str] = None,
+) -> tuple[str, str]:
     """
     Generate client certificate and key using mcp_security_framework.
     Args:
@@ -185,10 +253,10 @@ def generate_client_certificate(output_dir: str, ca_cert: str, client_name: str
         try:
             # Find CA key file
             ca_key = None
-            if ca_cert.endswith('.crt'):
-                ca_key = ca_cert.replace('.crt', '.key')
-            elif ca_cert.endswith('.pem'):
-                ca_key = ca_cert.replace('.pem', '_key.pem')
+            if ca_cert.endswith(".crt"):
+                ca_key = ca_cert.replace(".crt", ".key")
+            elif ca_cert.endswith(".pem"):
+                ca_key = ca_cert.replace(".pem", "_key.pem")
             if not os.path.exists(ca_key):
                 print(f"❌ CA key file not found: {ca_key}")
                 return None, None
@@ -206,7 +274,7 @@ def generate_client_certificate(output_dir: str, ca_cert: str, client_name: str
                 roles=roles or [],
                 permissions=permissions or [],
                 ca_cert_path=ca_cert,
-                ca_key_path=ca_key
+                ca_key_path=ca_key,
             )
             # Create certificate manager
             cert_config = CertificateConfig(
@@ -214,19 +282,25 @@ def generate_client_certificate(output_dir: str, ca_cert: str, client_name: str
                 key_storage_path=keys_dir,
                 default_validity_days=730,
                 key_size=2048,
-                hash_algorithm="sha256"
+                hash_algorithm="sha256",
             )
             cert_manager = CertificateManager(cert_config)
             # Create client certificate
             cert_pair = cert_manager.create_client_certificate(client_config)
             if cert_pair and cert_pair.certificate_path and cert_pair.private_key_path:
-                print(f"✅ Generated client certificate {client_name} using mcp_security_framework")
+                print(
+                    f"✅ Generated client certificate {client_name} using mcp_security_framework"
+                )
                 return (cert_pair.certificate_path, cert_pair.private_key_path)
             else:
-                print(f"❌ Failed to create client certificate {client_name}: Invalid certificate pair")
+                print(
+                    f"❌ Failed to create client certificate {client_name}: Invalid certificate pair"
+                )
                 return None, None
         except Exception as e:
-            print(f"❌ Error creating client certificate {client_name} with framework: {e}")
+            print(
+                f"❌ Error creating client certificate {client_name} with framework: {e}"
+            )
             return None, None
     else:
         # Fallback to OpenSSL
@@ -235,23 +309,56 @@ def generate_client_certificate(output_dir: str, ca_cert: str, client_name: str
         client_csr = os.path.join(certs_dir, f"{client_name}.csr")
         client_cert = os.path.join(certs_dir, f"{client_name}.crt")
         # Generate client private key
-        subprocess.run([
-            "openssl", "genrsa", "-out", client_key, "2048"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            ["openssl", "genrsa", "-out", client_key, "2048"],
+            check=True,
+            capture_output=True,
+        )
         # Generate client certificate signing request
-        subprocess.run([
-            "openssl", "req", "-new", "-key", client_key, "-out", client_csr,
-            "-subj", f"/C=US/ST=State/L=City/O=Organization/CN={client_name}-client"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            [
+                "openssl",
+                "req",
+                "-new",
+                "-key",
+                client_key,
+                "-out",
+                client_csr,
+                "-subj",
+                f"/C=US/ST=State/L=City/O=Organization/CN={client_name}-client",
+            ],
+            check=True,
+            capture_output=True,
+        )
         # Sign client certificate with CA
-        subprocess.run([
-            "openssl", "x509", "-req", "-in", client_csr, "-CA", ca_cert,
-            "-CAkey", ca_key, "-CAcreateserial", "-out", client_cert, "-days", "730"
-        ], check=True, capture_output=True)
+        subprocess.run(
+            [
+                "openssl",
+                "x509",
+                "-req",
+                "-in",
+                client_csr,
+                "-CA",
+                ca_cert,
+                "-CAkey",
+                ca_key,
+                "-CAcreateserial",
+                "-out",
+                client_cert,
+                "-days",
+                "730",
+            ],
+            check=True,
+            capture_output=True,
+        )
         # Clean up CSR
         os.remove(client_csr)
-        print(f"✅ Generated client certificate {client_name} using OpenSSL: {client_cert}")
+        print(
+            f"✅ Generated client certificate {client_name} using OpenSSL: {client_cert}"
+        )
         return client_cert, client_key
+
+
 def generate_tokens(output_dir: str) -> Dict[str, str]:
     """
     Generate API tokens for different roles.
@@ -267,14 +374,16 @@ def generate_tokens(output_dir: str) -> Dict[str, str]:
         "user": "user-token-456",
         "readonly": "readonly-token-123",
         "guest": "guest-token-123",
-        "proxy": "proxy-token-123"
+        "proxy": "proxy-token-123",
     }
     # Save tokens to file
     tokens_file = os.path.join(tokens_dir, "tokens.json")
-    with open(tokens_file, 'w') as f:
+    with open(tokens_file, "w") as f:
         json.dump(tokens, f, indent=2)
     print(f"✅ Generated tokens: {tokens_file}")
     return tokens
+
+
 def generate_roles_config(output_dir: str) -> Dict[str, Any]:
     """
     Generate roles configuration file.
@@ -285,37 +394,54 @@ def generate_roles_config(output_dir: str) -> Dict[str, Any]:
     """
     roles_config = {
         "admin": {
-            "permissions": ["read", "write", "execute", "delete", "admin", "register", "unregister", "heartbeat", "discover"],
-            "tokens": []
+            "permissions": [
+                "read",
+                "write",
+                "execute",
+                "delete",
+                "admin",
+                "register",
+                "unregister",
+                "heartbeat",
+                "discover",
+            ],
+            "tokens": [],
         },
         "user": {
-            "permissions": ["read", "execute", "register", "unregister", "heartbeat", "discover"],
-            "tokens": []
-        },
-        "readonly": {
-            "permissions": ["read", "discover"],
-            "tokens": []
-        },
-        "guest": {
-            "permissions": ["read", "discover"],
-            "tokens": []
+            "permissions": [
+                "read",
+                "execute",
+                "register",
+                "unregister",
+                "heartbeat",
+                "discover",
+            ],
+            "tokens": [],
         },
+        "readonly": {"permissions": ["read", "discover"], "tokens": []},
+        "guest": {"permissions": ["read", "discover"], "tokens": []},
         "proxy": {
             "permissions": ["register", "unregister", "heartbeat", "discover"],
-            "tokens": []
-        }
+            "tokens": [],
+        },
     }
     # Save roles config to file
     roles_file = os.path.join(output_dir, "roles.json")
-    with open(roles_file, 'w') as f:
+    with open(roles_file, "w") as f:
         json.dump(roles_config, f, indent=2)
     print(f"✅ Generated roles configuration: {roles_file}")
     return roles_config
+
+
 def main():
     """Main function for certificate and token generation."""
     parser = argparse.ArgumentParser(description="Generate certificates and tokens")
-    parser.add_argument("--output-dir", "-o", default="./certs", help="Output directory")
-    parser.add_argument("--framework", action="store_true", help="Use mcp_security_framework")
+    parser.add_argument(
+        "--output-dir", "-o", default="./certs", help="Output directory"
+    )
+    parser.add_argument(
+        "--framework", action="store_true", help="Use mcp_security_framework"
+    )
     args = parser.parse_args()
     print("🔐 Certificate and Token Generation Script")
     print("=" * 50)
@@ -340,11 +466,29 @@ def main():
         # 3. Generate client certificates
         print("\n🔧 Generating client certificates...")
         client_configs = [
-            ("admin", ["admin"], ["read", "write", "execute", "delete", "admin", "register", "unregister", "heartbeat", "discover"]),
-            ("user", ["user"], ["read", "execute", "register", "unregister", "heartbeat", "discover"]),
+            (
+                "admin",
+                ["admin"],
+                [
+                    "read",
+                    "write",
+                    "execute",
+                    "delete",
+                    "admin",
+                    "register",
+                    "unregister",
+                    "heartbeat",
+                    "discover",
+                ],
+            ),
+            (
+                "user",
+                ["user"],
+                ["read", "execute", "register", "unregister", "heartbeat", "discover"],
+            ),
             ("readonly", ["readonly"], ["read", "discover"]),
             ("guest", ["guest"], ["read", "discover"]),
-            ("proxy", ["proxy"], ["register", "unregister", "heartbeat", "discover"])
+            ("proxy", ["proxy"], ["register", "unregister", "heartbeat", "discover"]),
         ]
         for client_name, roles, permissions in client_configs:
             client_cert, client_key = generate_client_certificate(
@@ -365,5 +509,7 @@ def main():
     except Exception as e:
         print(f"❌ Error during generation: {e}")
         return 1
+
+
 if __name__ == "__main__":
     exit(main())