Load yaml safely

disk-usage-warn

@@ -1,4 +1,4 @@
-#!/usr/env python3
+#!/usr/bin/env python3
 
 """
 
@@ -14,6 +14,7 @@ Released under the GNU GENERAL PUBLIC LICENSE v3 (See LICENSE file for more)
 
 
 #
+import os
 import re
 import subprocess
 import sys
@@ -78,8 +79,27 @@ class DiskUsageWarn:
 		self.__configs = []
 		
 		#
-		for config_path in self.__config_paths:
-			self.consume_config(config_path)
+		for path in self.__config_paths:
+			
+			if os.path.isdir(path):
+				
+				self.log(path + " is actually a directory; Iterating over contained files (non-recursive)")
+				for file_name in os.listdir(path):
+					
+					one_config_path = os.path.join(path, file_name)
+					if re.match(".+\.yaml$", file_name):
+						self.log("Found yaml: " + file_name)
+						self.consume_config(one_config_path)
+					else:
+						self.log("Ignoring non-yaml file: " + file_name)
+			
+			elif os.path.isfile(path):
+				
+				self.consume_config(path)
+				
+			else:
+				
+				raise Exception("Don't know what to do with config path:" + str(path))
 		
 		self.log("Consumed " + str(len(self.__configs)) + " configs")
 		
@@ -99,7 +119,7 @@ class DiskUsageWarn:
 			raise Exception("Unable to open config file: " + path)
 		
 		# Parse
-		config = yaml.load(f)
+		config = yaml.safe_load(f)
 		
 		# Add the config file's own path
 		config["path"] = path
@@ -108,8 +128,6 @@ class DiskUsageWarn:
 	
 	def run(self):
 		
-		self.log("Begin")
-		
 		self.consume_configs()
 		self.do_configs()
 	
@@ -121,15 +139,15 @@ class DiskUsageWarn:
 	def do_config(self, config):
 		
 		# Pull the max usage
-		if "max-usage" not in config.keys():
-			raise Exception("Did not find config key: max-usage")
-		max_percent = config["max-usage"]
+		self.assert_config_key(config, "max-usage", [int, str, list])
+		max_percent = str(config["max-usage"])
 		match = re.match("(?P<integer_percent>[0-9]+)%?", max_percent)
 		if not match:
 			raise Exception("Unable to parse configuration value for max-usage (integer percent)")
 		max_percent = int(match.group("integer_percent"))
 		
 		# Check each device
+		self.assert_config_key(config, "devices", list)
 		for device in config["devices"]:
 			
 			device_usage = self.get_device_usage(device)
@@ -141,6 +159,23 @@ class DiskUsageWarn:
 				)
 				self.stderr(error_message)
 	
+	@staticmethod
+	def assert_config_key(config, key, types):
+		
+		if not isinstance(types, list):
+			types = [types]
+		if len(types) == 0:
+			types = [list]
+		
+		if key not in config.keys():
+			raise Exception("Missing required config key: " + str(key))
+		
+		for t in types:
+			if isinstance(config[key], t):
+				return True
+		
+		raise Exception("Config key \"" + key + "\" should be one of these types \"" + str(types) + "\"")
+	
 	def get_device_usage(self, device):
 		
 		args = ["df", device]