Skip to content

Commit

Permalink
test: add SSL test
Browse files Browse the repository at this point in the history
This adds a simple SSL test along with the framework for running
the test Zookeeper in a mode where it listens on both SSL and
non-SSL ports.
  • Loading branch information
James E. Blair authored and StephenSorriaux committed Jan 23, 2024
1 parent a440c91 commit be1b772
Show file tree
Hide file tree
Showing 4 changed files with 196 additions and 3 deletions.
155 changes: 153 additions & 2 deletions kazoo/testing/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@
import tempfile
import traceback

import OpenSSL
import jks


log = logging.getLogger(__name__)

Expand Down Expand Up @@ -67,7 +70,8 @@ def to_java_compatible_path(path):

ServerInfo = namedtuple(
"ServerInfo",
"server_id client_port election_port leader_port admin_port peer_type",
"server_id client_port secure_client_port "
"election_port leader_port admin_port peer_type",
)


Expand All @@ -88,6 +92,7 @@ def __init__(
configuration_entries=(),
java_system_properties=(),
jaas_config=None,
ssl_configuration=None,
):
"""Define the ZooKeeper test instance.
Expand All @@ -104,6 +109,9 @@ def __init__(
self.configuration_entries = configuration_entries
self.java_system_properties = java_system_properties
self.jaas_config = jaas_config
self.ssl_configuration = (
ssl_configuration if ssl_configuration is not None else {}
)

def run(self):
"""Run the ZooKeeper instance under a temporary directory.
Expand All @@ -117,6 +125,8 @@ def run(self):
log_path = os.path.join(self.working_path, "log")
log4j_path = os.path.join(self.working_path, "log4j.properties")
data_path = os.path.join(self.working_path, "data")
truststore_path = os.path.join(self.working_path, "truststore.jks")
keystore_path = os.path.join(self.working_path, "keystore.jks")

# various setup steps
if not os.path.exists(self.working_path):
Expand All @@ -126,21 +136,39 @@ def run(self):
if not os.path.exists(data_path):
os.mkdir(data_path)

try:
self.ssl_configuration["truststore"].save(
truststore_path, "apassword"
)
self.ssl_configuration["keystore"].save(keystore_path, "apassword")
except Exception:
log.exception("Unable to perform SSL configuration: ")
raise

with open(config_path, "w") as config:
config.write(
"""
tickTime=2000
dataDir=%s
clientPort=%s
secureClientPort=%s
maxClientCnxns=0
admin.serverPort=%s
serverCnxnFactory=org.apache.zookeeper.server.NettyServerCnxnFactory
authProvider.1=org.apache.zookeeper.server.auth.SASLAuthenticationProvider
ssl.keyStore.location=%s
ssl.keyStore.password=apassword
ssl.trustStore.location=%s
ssl.trustStore.password=apassword
%s
"""
% (
to_java_compatible_path(data_path),
self.server_info.client_port,
self.server_info.secure_client_port,
self.server_info.admin_port,
to_java_compatible_path(keystore_path),
to_java_compatible_path(truststore_path),
"\n".join(self.configuration_entries),
)
) # NOQA
Expand Down Expand Up @@ -266,6 +294,11 @@ def address(self):
"""Get the address of the ZooKeeper instance."""
return "%s:%s" % (self.host, self.client_port)

@property
def secure_address(self):
"""Get the address of the SSL ZooKeeper instance."""
return "%s:%s" % (self.host, self.secure_client_port)

@property
def running(self):
return self._running
Expand All @@ -274,6 +307,10 @@ def running(self):
def client_port(self):
return self.server_info.client_port

@property
def secure_client_port(self):
return self.server_info.secure_client_port

def reset(self):
"""Stop the zookeeper instance, cleaning out its on disk-data."""
self.stop()
Expand Down Expand Up @@ -329,6 +366,8 @@ def __init__(
self._install_path = install_path
self._classpath = classpath
self._servers = []
self._ssl_configuration = {}
self.perform_ssl_certs_generation()

# Calculate ports and peer group
port = port_offset
Expand All @@ -341,7 +380,13 @@ def __init__(
else:
peer_type = "participant"
info = ServerInfo(
server_id, port, port + 1, port + 2, port + 3, peer_type
server_id,
port,
port + 4,
port + 1,
port + 2,
port + 3,
peer_type,
)
peers.append(info)
port += 10
Expand All @@ -359,6 +404,7 @@ def __init__(
configuration_entries=configuration_entries,
java_system_properties=java_system_properties,
jaas_config=jaas_config,
ssl_configuration=dict(self._ssl_configuration),
)
)

Expand Down Expand Up @@ -399,3 +445,108 @@ def get_logs(self):
for server in self:
logs += server.get_logs()
return logs

def perform_ssl_certs_generation(self):
if self._ssl_configuration:
return

# generate CA key
ca_key = OpenSSL.crypto.PKey()
ca_key.generate_key(OpenSSL.crypto.TYPE_RSA, 2048)

# generate CA
ca_cert = OpenSSL.crypto.X509()
ca_cert.set_version(2)
ca_cert.set_serial_number(1)
ca_cert.get_subject().CN = "ca.kazoo.org"
ca_cert.gmtime_adj_notBefore(0)
ca_cert.gmtime_adj_notAfter(24 * 60 * 60)
ca_cert.set_issuer(ca_cert.get_subject())
ca_cert.set_pubkey(ca_key)
ca_cert.add_extensions(
[
OpenSSL.crypto.X509Extension(
b"basicConstraints", True, b"CA:TRUE, pathlen:0"
),
OpenSSL.crypto.X509Extension(
b"keyUsage", True, b"keyCertSign, cRLSign"
),
OpenSSL.crypto.X509Extension(
b"subjectKeyIdentifier", False, b"hash", subject=ca_cert
),
]
)
ca_cert.sign(ca_key, "sha256")

# generate server cert
server_key = OpenSSL.crypto.PKey()
server_key.generate_key(OpenSSL.crypto.TYPE_RSA, 2048)
server_cert = OpenSSL.crypto.X509()
server_cert.get_subject().CN = "localhost"
server_cert.set_serial_number(2)
server_cert.gmtime_adj_notBefore(0)
server_cert.gmtime_adj_notAfter(24 * 60 * 60)
server_cert.set_issuer(ca_cert.get_subject())
server_cert.set_pubkey(server_key)
server_cert.sign(ca_key, "sha256")

# generate client cert
client_key = OpenSSL.crypto.PKey()
client_key.generate_key(OpenSSL.crypto.TYPE_RSA, 2048)
client_cert = OpenSSL.crypto.X509()
client_cert.get_subject().CN = "client"
client_cert.set_serial_number(3)
client_cert.gmtime_adj_notBefore(0)
client_cert.gmtime_adj_notAfter(24 * 60 * 60)
client_cert.set_issuer(ca_cert.get_subject())
client_cert.set_pubkey(client_key)
client_cert.sign(ca_key, "sha256")

dumped_ca_cert = OpenSSL.crypto.dump_certificate(
OpenSSL.crypto.FILETYPE_ASN1, ca_cert
)

tce = jks.TrustedCertEntry.new("kazoo ca", dumped_ca_cert)
truststore = jks.KeyStore.new("jks", [tce])

dumped_server_cert = OpenSSL.crypto.dump_certificate(
OpenSSL.crypto.FILETYPE_ASN1, server_cert
)
dumped_server_key = OpenSSL.crypto.dump_privatekey(
OpenSSL.crypto.FILETYPE_ASN1, server_key
)

server_pke = jks.PrivateKeyEntry.new(
"server cert", [dumped_server_cert], dumped_server_key, "rsa_raw"
)

keystore = jks.KeyStore.new("jks", [server_pke])

self._ssl_configuration = {
"ca_cert": ca_cert,
"ca_key": ca_key,
"ca_cert_pem": OpenSSL.crypto.dump_certificate(
OpenSSL.crypto.FILETYPE_PEM, ca_cert
),
"server_cert": server_cert,
"server_key": server_key,
"client_cert": client_cert,
"client_key": client_key,
"client_cert_pem": OpenSSL.crypto.dump_certificate(
OpenSSL.crypto.FILETYPE_PEM, client_cert
),
"client_key_pem": OpenSSL.crypto.dump_privatekey(
OpenSSL.crypto.FILETYPE_PEM, client_key
),
"truststore": truststore,
"keystore": keystore,
}

def get_ssl_client_configuration(self):
if not self._ssl_configuration:
raise RuntimeError("SSL not configured yet.")
return {
"client_key": self._ssl_configuration["client_key_pem"],
"client_cert": self._ssl_configuration["client_cert_pem"],
"ca_cert": self._ssl_configuration["ca_cert_pem"],
}
9 changes: 8 additions & 1 deletion kazoo/testing/harness.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,10 @@ def log(self, level, msg, *args, **kwargs):
def servers(self):
return ",".join([s.address for s in self.cluster])

@property
def secure_servers(self):
return ",".join([s.secure_address for s in self.cluster])

def _get_nonchroot_client(self):
c = KazooClient(self.servers)
self._clients.append(c)
Expand Down Expand Up @@ -234,7 +238,10 @@ def setup_zookeeper(self, **client_options):
self.cluster.terminate()
self.cluster.start()
continue

if client_options.get("use_ssl"):
self.hosts = self.secure_servers + namespace
else:
self.hosts = self.servers + namespace
self.client = self._get_client(**client_options)
self.client.start()
self.client.ensure_path("/")
Expand Down
33 changes: 33 additions & 0 deletions kazoo/tests/test_client.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import os
import socket
import sys
import tempfile
import threading
import time
import uuid
Expand Down Expand Up @@ -1159,6 +1161,37 @@ def test_request_queuing_session_expired(self):
client.stop()


class TestSSLClient(KazooTestCase):
def setUp(self):
if CI_ZK_VERSION and CI_ZK_VERSION < (3, 5):
pytest.skip("Must use Zookeeper 3.5 or above")
ssl_path = tempfile.mkdtemp()
key_path = os.path.join(ssl_path, "key.pem")
cert_path = os.path.join(ssl_path, "cert.pem")
cacert_path = os.path.join(ssl_path, "cacert.pem")
with open(key_path, "wb") as key_file:
key_file.write(
self.cluster.get_ssl_client_configuration()["client_key"]
)
with open(cert_path, "wb") as cert_file:
cert_file.write(
self.cluster.get_ssl_client_configuration()["client_cert"]
)
with open(cacert_path, "wb") as cacert_file:
cacert_file.write(
self.cluster.get_ssl_client_configuration()["ca_cert"]
)
self.setup_zookeeper(
use_ssl=True, keyfile=key_path, certfile=cert_path, ca=cacert_path
)

def test_create(self):
client = self.client
path = client.create("/1")
assert path == "/1"
assert client.exists("/1")


dummy_dict = {
"aversion": 1,
"ctime": 0,
Expand Down
2 changes: 2 additions & 0 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ test =
pytest-cov
gevent>=1.2 ; implementation_name!='pypy'
eventlet>=0.17.1 ; implementation_name!='pypy'
pyjks
pyopenssl

eventlet =
eventlet>=0.17.1
Expand Down

0 comments on commit be1b772

Please sign in to comment.