From d473d237256583b57449498c731dc2183eb45f5b Mon Sep 17 00:00:00 2001
From: Jarrod Johnson <jjohnson2@lenovo.com>
Date: Thu, 22 Feb 2018 09:55:46 -0500
Subject: [PATCH] Add ability to specify attribs from environment

This is of particular use for passing passwords safely
into nodeattrib/nodegroupattrib.
---
 confluent_client/bin/nodeattrib          |  6 +++++-
 confluent_client/bin/nodegroupattrib     |  8 ++++++--
 confluent_client/confluent/client.py     | 19 +++++++++++++++++++
 confluent_client/doc/man/nodeattrib.ronn |  3 +++
 4 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/confluent_client/bin/nodeattrib b/confluent_client/bin/nodeattrib
index 7c3c6c6c..0f5e9cd6 100755
--- a/confluent_client/bin/nodeattrib
+++ b/confluent_client/bin/nodeattrib
@@ -37,10 +37,14 @@ import confluent.client as client
 argparser = optparse.OptionParser(
     usage='''\n       %prog [-b] noderange [list of attributes]  \
              \n       %prog -c noderange <list of attributes>  \
+             \n       %prog -e noderange <attribute names to set>  \
              \n       %prog noderange attribute1=value1 attribute2=value,...
              \n ''')
 argparser.add_option('-b', '--blame', action='store_true',
                      help='Show information about how attributes inherited')
+argparser.add_option('-e', '--environment', action='store_true',
+                     help='Set attributes, but from environment variable of '
+                          'same name')
 argparser.add_option('-c', '--clear', action='store_true',
                      help='Clear attributes')
 (options, args) = argparser.parse_args()
@@ -63,7 +67,7 @@ exitcode = 0
 nodetype="noderange"
 
 if len(args) > 1:
-    if "=" in args[1] or options.clear:
+    if "=" in args[1] or options.clear or options.environment:
         if "=" in args[1] and options.clear:
             print("Can not clear and set at the same time!")
             argparser.print_help()
diff --git a/confluent_client/bin/nodegroupattrib b/confluent_client/bin/nodegroupattrib
index c695ad34..29d95752 100755
--- a/confluent_client/bin/nodegroupattrib
+++ b/confluent_client/bin/nodegroupattrib
@@ -36,11 +36,15 @@ import confluent.client as client
 argparser = optparse.OptionParser(
     usage='''\n       %prog [options]   \
              \n       %prog [options] nodegroup [list of attributes]  \
-             \n       %prog [options] nodegroup group=value1,value2   \
-             \n       %prog [options] nodegroup group=value1,value2
+             \n       %prog [options] nodegroup nodes=value1,value2   \
+             \n       %prog -e nodegroup <attribute names to set>  \             
+             \n       %prog [options] nodegroup nodes=value1,value2
              \n ''')
 argparser.add_option('-b', '--blame', action='store_true',
                      help='Show information about how attributes inherited')
+argparser.add_option('-e', '--environment', action='store_true',
+                     help='Set attributes, but from environment variable of '
+                          'same name')
 argparser.add_option('-c', '--clear', action='store_true',
                      help='Clear variables')
 (options, args) = argparser.parse_args()
diff --git a/confluent_client/confluent/client.py b/confluent_client/confluent/client.py
index 92607f69..aa4b2224 100644
--- a/confluent_client/confluent/client.py
+++ b/confluent_client/confluent/client.py
@@ -480,6 +480,25 @@ def updateattrib(session, updateargs, nodetype, noderange, options):
                     exitcode = res['errorcode']
                 sys.stderr.write('Error: ' + res['error'] + '\n')
         sys.exit(exitcode)
+    elif hasattr(options, 'environment') and options.environment:
+        for key in updateargs[1:]:
+            key = key.replace('.', '_')
+            value = os.environ.get(
+                key, os.environ[key.upper()])
+            # Let's do one pass to make sure that there's not a usage problem
+        for key in updateargs[1:]:
+            key = key.replace('.', '_')
+            value = os.environ.get(
+                key, os.environ[key.upper()])
+            if (nodetype == "nodegroups"):
+                exitcode = session.simple_nodegroups_command(noderange,
+                                                             'attributes/all',
+                                                             value, key)
+            else:
+                exitcode = session.simple_noderange_command(noderange,
+                                                            'attributes/all',
+                                                            value, key)
+        sys.exit(exitcode)
     else:
         if "=" in updateargs[1]:
             try:
diff --git a/confluent_client/doc/man/nodeattrib.ronn b/confluent_client/doc/man/nodeattrib.ronn
index ca3f830a..03601c72 100644
--- a/confluent_client/doc/man/nodeattrib.ronn
+++ b/confluent_client/doc/man/nodeattrib.ronn
@@ -6,6 +6,7 @@ nodeattrib(8) -- List or change confluent nodes attributes
 `nodeattrib [-b] <noderange> [<nodeattribute>...]`  
 `nodeattrib <noderange> [<nodeattribute1=value1> <nodeattribute2=value2> ...]`  
 `nodeattrib -c <noderange> <nodeattribute1> <nodeattribute2> ...`  
+`nodeattrib -e <noderange> <nodeattribute1> <nodeattribute2> ...`  
 
 ## DESCRIPTION
 
@@ -33,6 +34,8 @@ See nodegroupattrib(8) command on how to manage attributes on a group level.
   Annotate inherited and expression based attributes to show their base value.
 * `-c`, `--clear`:
   Clear specified nodeattributes
+* `-e`, `--environment`:
+  Set specified attributes based on exported environment variable of matching name.  Environment variable names may be lower case or all upper case.  Replace . with _ as needed (e.g. info.note may be specified as either $info_note or $INFO_NOTE
 
 ## EXAMPLES
 * Listing matching nodes of a simple noderange: