From ee66ce8388b70f4c3cbefa6cddf34437848f8393 Mon Sep 17 00:00:00 2001
From: Anton Ivashkin <iantonspb@yandex-team.ru>
Date: Wed, 20 Dec 2023 19:05:56 +0200
Subject: [PATCH] Add '--secure' flag in 'chadmin zookeeper' commands

---
 ch_tools/chadmin/cli/zookeeper_group.py | 6 +++++-
 ch_tools/chadmin/internal/zookeeper.py  | 9 ++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/ch_tools/chadmin/cli/zookeeper_group.py b/ch_tools/chadmin/cli/zookeeper_group.py
index 154106c5..2c7713ec 100644
--- a/ch_tools/chadmin/cli/zookeeper_group.py
+++ b/ch_tools/chadmin/cli/zookeeper_group.py
@@ -23,6 +23,8 @@
 @group("zookeeper")
 @option("--port", help="ZooKeeper port.", type=int, default=2181)
 @option("--host", help="ZooKeeper host.", type=str)
+@option("--secure", help="Use secure connection.", default=False, is_flag=True)
+@option("--no-verify-ssl-certs", help="Do not check SSL Certificates in secure connection.", default=False, is_flag=True)
 @option("--timeout", help="ZooKeeper timeout.", default=10)
 @option(
     "--zkcli-identity",
@@ -52,7 +54,7 @@
 )
 @pass_context
 def zookeeper_group(
-    ctx, host, port, timeout, zkcli_identity, no_chroot, no_ch_config, zk_root_path
+    ctx, host, secure, no_verify_ssl_certs, port, timeout, zkcli_identity, no_chroot, no_ch_config, zk_root_path
 ):
     """ZooKeeper management commands.
 
@@ -64,6 +66,8 @@ def zookeeper_group(
     ctx.obj["zk_client_args"] = {
         "port": port,
         "host": host,
+        "use_ssl": secure,
+        "no_verify_ssl_certs": no_verify_ssl_certs,
         "timeout": timeout,
         "zkcli_identity": zkcli_identity,
         "no_chroot": no_chroot,
diff --git a/ch_tools/chadmin/internal/zookeeper.py b/ch_tools/chadmin/internal/zookeeper.py
index 9116128b..4ff6cf83 100644
--- a/ch_tools/chadmin/internal/zookeeper.py
+++ b/ch_tools/chadmin/internal/zookeeper.py
@@ -303,6 +303,8 @@ def _get_zk_client(ctx):
     args = ctx.obj.get("zk_client_args", {})
     host = args.get("host")
     port = args.get("port", 2181)
+    use_ssl = args.get("use_ssl", False)
+    no_verify_ssl_certs = args.get("no_verify_ssl_certs", False)
     timeout = args.get("timeout", 10)
     zkcli_identity = args.get("zkcli_identity")
     no_chroot = args.get("no_chroot", False)
@@ -334,5 +336,10 @@ def _get_zk_client(ctx):
         auth_data = [("digest", zkcli_identity)]
 
     return KazooClient(
-        connect_str, auth_data=auth_data, timeout=timeout, logger=logging.getLogger()
+        connect_str,
+        auth_data=auth_data,
+        timeout=timeout,
+        logger=logging.getLogger(),
+        use_ssl=use_ssl,
+        verify_certs=not no_verify_ssl_certs,
     )