Improve node classifier.

1 files changed, 23 insertions, 2 deletions

diff --git a/files/node-classifier.py b/files/node-classifier.py
index 8c28ddb..7fc096f 100644
--- a/files/node-classifier.py
+++ b/files/node-classifier.py
@@ -2,17 +2,38 @@
 
 import yaml
 import sys
+import configparser
 
 # https://puppet.com/docs/puppet/7/nodes_external.html
 
-with open('/nodes.yaml') as f:
-    data = yaml.load(f)
 
 def main(args):
     if len(args) == 1:
         print('usage ./node-classifier <node-name>')
         return
 
+    parser = configparser.ConfigParser()
+    conf_file = '/etc/node-classifier.ini'
+    if parser.read(conf_file) != [conf_file]:
+        print(f'Configuration file missing, expected {conf_file}')
+        return 1
+
+    try:
+        fmt = parser.get('common', 'node_fmt')
+        if fmt != 'yaml':
+            print(f'Unknown format {fmt}')
+            return 1
+        filename = parser.get('common', 'nodes')
+    except configparser.NoSectionError as e:
+        print(e)
+        return 1
+    except NoOptionError as e:
+        print(e)
+        return 1
+
+    with open(filename) as f:
+        data = yaml.load(f)
+
     nodename = sys.argv[1]
     instance = data.get(nodename)
     if not instance: