From 27aa23a61d14ac747499e490294f181402bf8419 Mon Sep 17 00:00:00 2001
From: Tobias Wolf <NotTheEvilOne@users.noreply.github.com>
Date: Mon, 6 May 2024 14:36:53 +0200
Subject: [PATCH] Add YAML config schema and validation

Signed-off-by: Tobias Wolf <wolf@b1-systems.de>
---
 .pre-commit-config.yaml        |  1 +
 requirements.txt               |  1 +
 setup.cfg                      |  3 +++
 src/rookify/__main__.py        | 11 +++++------
 src/rookify/config.schema.yaml | 26 ++++++++++++++++++++++++++
 src/rookify/yaml.py            | 26 ++++++++++++++++++++++++--
 6 files changed, 60 insertions(+), 8 deletions(-)
 create mode 100644 src/rookify/config.schema.yaml

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index baeef2d..cf2f956 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -9,6 +9,7 @@ repos:
     rev: v4.5.0
     hooks:
       - id: check-yaml
+        args: [ --allow-multiple-documents ]
       - id: check-json
       - id: check-toml
       - id: check-merge-conflict
diff --git a/requirements.txt b/requirements.txt
index 219a759..9246059 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -27,5 +27,6 @@ rsa==4.9
 six==1.16.0
 structlog==24.1.0
 urllib3==2.2.1
+yamale==5.1.0
 websocket-client==1.7.0
 wrapt==1.16.0
diff --git a/setup.cfg b/setup.cfg
index d1edfe8..0b4de7f 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -7,3 +7,6 @@ install_requires=file:requirements.txt
 [options.extras_require]
 tests =
     pytest==8.0.2
+
+[options.package_data]
+rookify=config.schema.yaml
diff --git a/src/rookify/__main__.py b/src/rookify/__main__.py
index 990d768..89c8853 100644
--- a/src/rookify/__main__.py
+++ b/src/rookify/__main__.py
@@ -2,16 +2,15 @@
 
 import os
 import rookify.modules
-
 from types import MappingProxyType
-from .yaml import load_yaml, save_yaml
-from rookify.logger import configure_logging, get_logger
+from .logger import configure_logging, get_logger
+from .yaml import load_config, load_module_data, save_module_data
 
 
 def main() -> None:
     # Load configuration file
     try:
-        config = load_yaml("config.yaml")
+        config = load_config("config.yaml")
     except FileNotFoundError as err:
         raise SystemExit(f"Could not load config: {err}")
 
@@ -31,7 +30,7 @@ def main() -> None:
     module_data = dict()
 
     try:
-        module_data.update(load_yaml(config["general"]["module_data_file"]))
+        module_data.update(load_module_data(config["general"]["module_data_file"]))
     except FileNotFoundError:
         pass
 
@@ -70,7 +69,7 @@ def main() -> None:
         result = handler.run()
         module_data[migration_module.MODULE_NAME] = result
 
-    save_yaml(config["general"]["module_data_file"], module_data)
+    save_module_data(config["general"]["module_data_file"], module_data)
 
     log.info("Data was updated to module_data_file.")
 
diff --git a/src/rookify/config.schema.yaml b/src/rookify/config.schema.yaml
new file mode 100644
index 0000000..7b09587
--- /dev/null
+++ b/src/rookify/config.schema.yaml
@@ -0,0 +1,26 @@
+general:
+  module_data_file: str()
+
+ceph:
+  conf_file: str()
+  keyring: str()
+
+ssh:
+  private_key: str()
+  hosts: map(include("ssh_host"), key=str(), min=1)
+
+kubernetes:
+  config: str()
+
+rook:
+  cluster:
+    name: str()
+    namespace: str()
+  ceph:
+    image: str()
+
+migration_modules: list(str())
+---
+ssh_host:
+    address: ip()
+    user: str()
diff --git a/src/rookify/yaml.py b/src/rookify/yaml.py
index 856716f..6712ae4 100644
--- a/src/rookify/yaml.py
+++ b/src/rookify/yaml.py
@@ -1,16 +1,38 @@
 # -*- coding: utf-8 -*-
 
+import importlib.resources
+import importlib.resources.abc
+import yamale
 import yaml
+from pathlib import Path
 from typing import Any, Dict
 
 
-def load_yaml(path: str) -> Dict[str, Any]:
+_config_schema_file: Path | importlib.resources.abc.Traversable = Path(
+    "rookify", "config.schema.yaml"
+)
+for entry in importlib.resources.files("rookify").iterdir():
+    if entry.name == "config.schema.yaml":
+        _config_schema_file = entry
+
+
+def load_config(path: str) -> Dict[str, Any]:
+    schema = yamale.make_schema(_config_schema_file)
+    data = yamale.make_data(path)
+
+    yamale.validate(schema, data)
+
+    assert isinstance(data[0][0], dict)
+    return data[0][0]
+
+
+def load_module_data(path: str) -> Dict[str, Any]:
     with open(path, "r") as file:
         data = yaml.safe_load(file)
         assert isinstance(data, dict)
         return data
 
 
-def save_yaml(path: str, data: Dict[str, Any]) -> None:
+def save_module_data(path: str, data: Dict[str, Any]) -> None:
     with open(path, "w") as file:
         yaml.safe_dump(data, file)