From a085dbceac4e911a743c84388b55a0a8f2ba320b Mon Sep 17 00:00:00 2001
From: Scott Gasch <scott@gasch.org>
Date: Wed, 17 Aug 2022 11:04:31 -0700
Subject: [PATCH] Allow config.py to save/load config from zookeeper.  No
 watching yet.

---
 config.py     | 77 +++++++++++++++++++++++++++++++++++++++++++++++----
 file_utils.py | 17 ++++++++++++
 2 files changed, 88 insertions(+), 6 deletions(-)

diff --git a/config.py b/config.py
index c7c686d..f3a2dfe 100644
--- a/config.py
+++ b/config.py
@@ -85,9 +85,14 @@ import argparse
 import logging
 import os
 import pprint
+import re
 import sys
 from typing import Any, Dict, List, Optional
 
+from kazoo.client import KazooClient
+
+import scott_secrets
+
 # This module is commonly used by others in here and should avoid
 # taking any unnecessary dependencies back on them.
 
@@ -336,10 +341,35 @@ def _augment_sys_argv_from_loadfile():
             saw_other_args = True
 
     if loadfile is not None:
-        if not os.path.exists(loadfile):
+        zkpath = None
+        if loadfile[:3] == 'zk:':
+            try:
+                zk = KazooClient(
+                    hosts=scott_secrets.ZOOKEEPER_NODES,
+                    use_ssl=True,
+                    verify_certs=False,
+                    keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                    keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
+                    certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                )
+                zk.start()
+                zkpath = loadfile[3:]
+                if not zkpath.startswith('/config/'):
+                    zkpath = '/config/' + zkpath
+                    zkpath = re.sub(r'//+', '/', zkpath)
+                if not zk.exists(zkpath):
+                    raise Exception(
+                        f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
+                    )
+            except Exception as e:
+                raise Exception(
+                    f'ERROR: Error talking with zookeeper while looking for {loadfile}'
+                ) from e
+        elif not os.path.exists(loadfile):
             raise Exception(
                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
             )
+
         if saw_other_args:
             msg = f'Augmenting commandline arguments with those from {loadfile}.'
         else:
@@ -347,9 +377,21 @@ def _augment_sys_argv_from_loadfile():
         print(msg, file=sys.stderr)
         SAVED_MESSAGES.append(msg)
 
-        with open(loadfile, 'r') as rf:
-            newargs = rf.readlines()
-        newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
+        newargs = []
+        if zkpath:
+            try:
+                contents = zk.get(zkpath)[0]
+                contents = contents.decode()
+                newargs = [
+                    arg.strip('\n') for arg in contents.split('\n') if 'config_savefile' not in arg
+                ]
+            except Exception as e:
+                raise Exception(f'Error reading {zkpath} from zookeeper.') from e
+            SAVED_MESSAGES.append(f'Loaded config from zookeeper from {zkpath}')
+        else:
+            with open(loadfile, 'r') as rf:
+                newargs = rf.readlines()
+            newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
         sys.argv += newargs
 
 
@@ -403,8 +445,31 @@ def parse(entry_module: Optional[str]) -> Dict[str, Any]:
     # Check for savefile and populate it if requested.
     savefile = config['config_savefile']
     if savefile and len(savefile) > 0:
-        with open(savefile, 'w') as wf:
-            wf.write("\n".join(ORIG_ARGV[1:]))
+        data = '\n'.join(ORIG_ARGV[1:])
+        if savefile[:3] == 'zk:':
+            zkpath = savefile[3:]
+            if not zkpath.startswith('/config/'):
+                zkpath = '/config/' + zkpath
+                zkpath = re.sub(r'//+', '/', zkpath)
+            try:
+                zk = KazooClient(
+                    hosts=scott_secrets.ZOOKEEPER_NODES,
+                    use_ssl=True,
+                    verify_certs=False,
+                    keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                    keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
+                    certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
+                )
+                zk.start()
+                if zk.exists(zkpath):
+                    zk.delete(zkpath)
+                zk.create(zkpath, data.encode())
+            except Exception as e:
+                raise Exception(f'Failed to create zookeeper path {zkpath}') from e
+            SAVED_MESSAGES.append(f'Saved config to zookeeper in {zkpath}')
+        else:
+            with open(savefile, 'w') as wf:
+                wf.write(data)
 
     # Also dump the config on stderr if requested.
     if config['config_dump']:
diff --git a/file_utils.py b/file_utils.py
index 7a64f9f..b7cd6bc 100644
--- a/file_utils.py
+++ b/file_utils.py
@@ -92,6 +92,23 @@ def remove(path: str) -> None:
     os.remove(path)
 
 
+def fix_multiple_slashes(path: str) -> str:
+    """Fixes multi-slashes in paths or path-like strings
+
+    Args:
+        path: the path in which to remove multiple slashes
+
+    >>> p = '/usr/local//etc/rc.d///file.txt'
+    >>> fix_multiple_slashes(p)
+    '/usr/local/etc/rc.d/file.txt'
+
+    >>> p = 'this is a test'
+    >>> fix_multiple_slashes(p) == p
+    True
+    """
+    return re.sub(r'/+', '/', path)
+
+
 def delete(path: str) -> None:
     """This is a convenience for my dumb ass who can't remember os.remove
     sometimes.
-- 
2.46.0