From a7026aba8a1f7df49f1a89bbdaebc1838e51984a Mon Sep 17 00:00:00 2001
From: "James E. Blair" <jim@acmegating.com>
Date: Wed, 9 Nov 2022 13:44:28 -0800
Subject: [PATCH] Add ssl support to zookeeper-statsd and fix latency handling

This adds optional SSL support to zookeeper-statsd.  This could
come in handy if we ever decide to turn off the plaintext
localhost-only port.

This also corrects the type handling for the latency value, which
can be a floating point.

Change-Id: Id39fc8bd924eda528723c40d2e7e24993a60d6a5
---
 docker/zookeeper-statsd/zookeeper-statsd.py | 35 ++++++++++++++++++---
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/docker/zookeeper-statsd/zookeeper-statsd.py b/docker/zookeeper-statsd/zookeeper-statsd.py
index abdcd705ba..fca38ce29f 100755
--- a/docker/zookeeper-statsd/zookeeper-statsd.py
+++ b/docker/zookeeper-statsd/zookeeper-statsd.py
@@ -21,6 +21,8 @@ import logging
 import re
 import socket
 import time
+import os
+import ssl
 
 from statsd.defaults.env import statsd
 
@@ -48,14 +50,23 @@ COUNTERS = [
 
 
 class Socket:
-    def __init__(self, host, port):
+    def __init__(self, host, port, ca_cert, client_cert, client_key):
         self.host = host
         self.port = port
+        self.ca_cert = ca_cert
+        self.client_cert = client_cert
+        self.client_key = client_key
         self.socket = None
 
     def open(self):
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         s.settimeout(5)
+        if self.client_key:
+            context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+            context.load_verify_locations(self.ca_cert)
+            context.load_cert_chain(self.client_cert, self.client_key)
+            context.check_hostname = False
+            s = context.wrap_socket(s, server_hostname=self.host)
         s.connect((self.host, self.port))
         self.socket = s
 
@@ -69,8 +80,13 @@ class Socket:
 
 
 class ZooKeeperStats:
-    def __init__(self, host, port=2181):
-        self.socket = Socket(host, port)
+    def __init__(self, host, port=None,
+                 ca_cert=None, client_cert=None, client_key=None):
+        if client_key:
+            port = port or 2281
+        else:
+            port = port or 2181
+        self.socket = Socket(host, port, ca_cert, client_cert, client_key)
         # The hostname to use when reporting stats (e.g., zk01)
         if host in ('localhost', '127.0.0.1', '::1'):
             self.hostname = socket.gethostname()
@@ -108,7 +124,11 @@ class ZooKeeperStats:
         base = 'zk.%s.' % (self.hostname,)
         for key in GAUGES:
             try:
-                value = int(stats.get(key, 0))
+                value = stats.get(key, '0')
+                if '.' in value:
+                    value = float(value)
+                else:
+                    value = int(value)
                 pipe.gauge(base + key, value)
             except Exception:
                 self.log.exception("Unable to process %s", key)
@@ -137,6 +157,11 @@ class ZooKeeperStats:
         self.reportStats(stats)
 
 
+ca_cert = os.environ.get("ZK_CA_CERT")
+client_cert = os.environ.get("ZK_CLIENT_CERT")
+client_key = os.environ.get("ZK_CLIENT_KEY")
+
 logging.basicConfig(level=logging.DEBUG)
-p = ZooKeeperStats('localhost')
+p = ZooKeeperStats('localhost', ca_cert=ca_cert,
+                   client_cert=client_cert, client_key=client_key)
 p.run()