3434import tempfile
3535import traceback
3636
37+ import OpenSSL
38+ import jks
39+
3740
3841log = logging .getLogger (__name__ )
3942
@@ -67,7 +70,8 @@ def to_java_compatible_path(path):
6770
6871ServerInfo = namedtuple (
6972 "ServerInfo" ,
70- "server_id client_port election_port leader_port admin_port peer_type" ,
73+ "server_id client_port secure_client_port "
74+ "election_port leader_port admin_port peer_type" ,
7175)
7276
7377
@@ -88,6 +92,7 @@ def __init__(
8892 configuration_entries = (),
8993 java_system_properties = (),
9094 jaas_config = None ,
95+ ssl_configuration = None ,
9196 ):
9297 """Define the ZooKeeper test instance.
9398
@@ -104,6 +109,9 @@ def __init__(
104109 self .configuration_entries = configuration_entries
105110 self .java_system_properties = java_system_properties
106111 self .jaas_config = jaas_config
112+ self .ssl_configuration = (
113+ ssl_configuration if ssl_configuration is not None else {}
114+ )
107115
108116 def run (self ):
109117 """Run the ZooKeeper instance under a temporary directory.
@@ -117,6 +125,8 @@ def run(self):
117125 log_path = os .path .join (self .working_path , "log" )
118126 log4j_path = os .path .join (self .working_path , "log4j.properties" )
119127 data_path = os .path .join (self .working_path , "data" )
128+ truststore_path = os .path .join (self .working_path , "truststore.jks" )
129+ keystore_path = os .path .join (self .working_path , "keystore.jks" )
120130
121131 # various setup steps
122132 if not os .path .exists (self .working_path ):
@@ -126,21 +136,39 @@ def run(self):
126136 if not os .path .exists (data_path ):
127137 os .mkdir (data_path )
128138
139+ try :
140+ self .ssl_configuration ["truststore" ].save (
141+ truststore_path , "apassword"
142+ )
143+ self .ssl_configuration ["keystore" ].save (keystore_path , "apassword" )
144+ except Exception :
145+ log .exception ("Unable to perform SSL configuration: " )
146+ raise
147+
129148 with open (config_path , "w" ) as config :
130149 config .write (
131150 """
132151tickTime=2000
133152dataDir=%s
134153clientPort=%s
154+ secureClientPort=%s
135155maxClientCnxns=0
136156admin.serverPort=%s
157+ serverCnxnFactory=org.apache.zookeeper.server.NettyServerCnxnFactory
137158authProvider.1=org.apache.zookeeper.server.auth.SASLAuthenticationProvider
159+ ssl.keyStore.location=%s
160+ ssl.keyStore.password=apassword
161+ ssl.trustStore.location=%s
162+ ssl.trustStore.password=apassword
138163%s
139164"""
140165 % (
141166 to_java_compatible_path (data_path ),
142167 self .server_info .client_port ,
168+ self .server_info .secure_client_port ,
143169 self .server_info .admin_port ,
170+ to_java_compatible_path (keystore_path ),
171+ to_java_compatible_path (truststore_path ),
144172 "\n " .join (self .configuration_entries ),
145173 )
146174 ) # NOQA
@@ -266,6 +294,11 @@ def address(self):
266294 """Get the address of the ZooKeeper instance."""
267295 return "%s:%s" % (self .host , self .client_port )
268296
297+ @property
298+ def secure_address (self ):
299+ """Get the address of the SSL ZooKeeper instance."""
300+ return "%s:%s" % (self .host , self .secure_client_port )
301+
269302 @property
270303 def running (self ):
271304 return self ._running
@@ -274,6 +307,10 @@ def running(self):
274307 def client_port (self ):
275308 return self .server_info .client_port
276309
310+ @property
311+ def secure_client_port (self ):
312+ return self .server_info .secure_client_port
313+
277314 def reset (self ):
278315 """Stop the zookeeper instance, cleaning out its on disk-data."""
279316 self .stop ()
@@ -329,6 +366,8 @@ def __init__(
329366 self ._install_path = install_path
330367 self ._classpath = classpath
331368 self ._servers = []
369+ self ._ssl_configuration = {}
370+ self .perform_ssl_certs_generation ()
332371
333372 # Calculate ports and peer group
334373 port = port_offset
@@ -341,7 +380,13 @@ def __init__(
341380 else :
342381 peer_type = "participant"
343382 info = ServerInfo (
344- server_id , port , port + 1 , port + 2 , port + 3 , peer_type
383+ server_id ,
384+ port ,
385+ port + 4 ,
386+ port + 1 ,
387+ port + 2 ,
388+ port + 3 ,
389+ peer_type ,
345390 )
346391 peers .append (info )
347392 port += 10
@@ -359,6 +404,7 @@ def __init__(
359404 configuration_entries = configuration_entries ,
360405 java_system_properties = java_system_properties ,
361406 jaas_config = jaas_config ,
407+ ssl_configuration = dict (self ._ssl_configuration ),
362408 )
363409 )
364410
@@ -399,3 +445,108 @@ def get_logs(self):
399445 for server in self :
400446 logs += server .get_logs ()
401447 return logs
448+
449+ def perform_ssl_certs_generation (self ):
450+ if self ._ssl_configuration :
451+ return
452+
453+ # generate CA key
454+ ca_key = OpenSSL .crypto .PKey ()
455+ ca_key .generate_key (OpenSSL .crypto .TYPE_RSA , 2048 )
456+
457+ # generate CA
458+ ca_cert = OpenSSL .crypto .X509 ()
459+ ca_cert .set_version (2 )
460+ ca_cert .set_serial_number (1 )
461+ ca_cert .get_subject ().CN = "ca.kazoo.org"
462+ ca_cert .gmtime_adj_notBefore (0 )
463+ ca_cert .gmtime_adj_notAfter (24 * 60 * 60 )
464+ ca_cert .set_issuer (ca_cert .get_subject ())
465+ ca_cert .set_pubkey (ca_key )
466+ ca_cert .add_extensions (
467+ [
468+ OpenSSL .crypto .X509Extension (
469+ b"basicConstraints" , True , b"CA:TRUE, pathlen:0"
470+ ),
471+ OpenSSL .crypto .X509Extension (
472+ b"keyUsage" , True , b"keyCertSign, cRLSign"
473+ ),
474+ OpenSSL .crypto .X509Extension (
475+ b"subjectKeyIdentifier" , False , b"hash" , subject = ca_cert
476+ ),
477+ ]
478+ )
479+ ca_cert .sign (ca_key , "sha256" )
480+
481+ # generate server cert
482+ server_key = OpenSSL .crypto .PKey ()
483+ server_key .generate_key (OpenSSL .crypto .TYPE_RSA , 2048 )
484+ server_cert = OpenSSL .crypto .X509 ()
485+ server_cert .get_subject ().CN = "localhost"
486+ server_cert .set_serial_number (2 )
487+ server_cert .gmtime_adj_notBefore (0 )
488+ server_cert .gmtime_adj_notAfter (24 * 60 * 60 )
489+ server_cert .set_issuer (ca_cert .get_subject ())
490+ server_cert .set_pubkey (server_key )
491+ server_cert .sign (ca_key , "sha256" )
492+
493+ # generate client cert
494+ client_key = OpenSSL .crypto .PKey ()
495+ client_key .generate_key (OpenSSL .crypto .TYPE_RSA , 2048 )
496+ client_cert = OpenSSL .crypto .X509 ()
497+ client_cert .get_subject ().CN = "client"
498+ client_cert .set_serial_number (3 )
499+ client_cert .gmtime_adj_notBefore (0 )
500+ client_cert .gmtime_adj_notAfter (24 * 60 * 60 )
501+ client_cert .set_issuer (ca_cert .get_subject ())
502+ client_cert .set_pubkey (client_key )
503+ client_cert .sign (ca_key , "sha256" )
504+
505+ dumped_ca_cert = OpenSSL .crypto .dump_certificate (
506+ OpenSSL .crypto .FILETYPE_ASN1 , ca_cert
507+ )
508+
509+ tce = jks .TrustedCertEntry .new ("kazoo ca" , dumped_ca_cert )
510+ truststore = jks .KeyStore .new ("jks" , [tce ])
511+
512+ dumped_server_cert = OpenSSL .crypto .dump_certificate (
513+ OpenSSL .crypto .FILETYPE_ASN1 , server_cert
514+ )
515+ dumped_server_key = OpenSSL .crypto .dump_privatekey (
516+ OpenSSL .crypto .FILETYPE_ASN1 , server_key
517+ )
518+
519+ server_pke = jks .PrivateKeyEntry .new (
520+ "server cert" , [dumped_server_cert ], dumped_server_key , "rsa_raw"
521+ )
522+
523+ keystore = jks .KeyStore .new ("jks" , [server_pke ])
524+
525+ self ._ssl_configuration = {
526+ "ca_cert" : ca_cert ,
527+ "ca_key" : ca_key ,
528+ "ca_cert_pem" : OpenSSL .crypto .dump_certificate (
529+ OpenSSL .crypto .FILETYPE_PEM , ca_cert
530+ ),
531+ "server_cert" : server_cert ,
532+ "server_key" : server_key ,
533+ "client_cert" : client_cert ,
534+ "client_key" : client_key ,
535+ "client_cert_pem" : OpenSSL .crypto .dump_certificate (
536+ OpenSSL .crypto .FILETYPE_PEM , client_cert
537+ ),
538+ "client_key_pem" : OpenSSL .crypto .dump_privatekey (
539+ OpenSSL .crypto .FILETYPE_PEM , client_key
540+ ),
541+ "truststore" : truststore ,
542+ "keystore" : keystore ,
543+ }
544+
545+ def get_ssl_client_configuration (self ):
546+ if not self ._ssl_configuration :
547+ raise RuntimeError ("SSL not configured yet." )
548+ return {
549+ "client_key" : self ._ssl_configuration ["client_key_pem" ],
550+ "client_cert" : self ._ssl_configuration ["client_cert_pem" ],
551+ "ca_cert" : self ._ssl_configuration ["ca_cert_pem" ],
552+ }
0 commit comments