Add base logic for grub lib (#98)

lib/charms/operator_libs_linux/v0/grub.py

@@ -1,24 +1,67 @@
-"""TODO: Add a proper docstring here.
+# Copyright 2023 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-This is a placeholder docstring for this charm library. Docstrings are
-presented on Charmhub and updated whenever you push a new version of the
-library.
+"""Simple library for managing Linux kernel configuration via GRUB.
 
-Complete documentation about creating and documenting libraries can be found
-in the SDK docs at https://juju.is/docs/sdk/libraries.
+This library is only used for setting additional parameters that will be stored in the
+"/etc/default/grub.d/95-juju-charm.cfg" config file and not for editing other
+configuration files. It's intended to be used in charms to help configure the machine.
 
-See `charmcraft publish-lib` and `charmcraft fetch-lib` for details of how to
-share and consume charm libraries. They serve to enhance collaboration
-between charmers. Use a charmer's libraries for classes that handle
-integration with their charm.
+Configurations for individual charms will be stored in "/etc/default/grub.d/90-juju-<charm>",
+but these configurations will only have informational value as all configurations will be merged
+to "/etc/default/grub.d/95-juju-charm.cfg".
 
-Bear in mind that new revisions of the different major API versions (v0, v1,
-v2 etc) are maintained independently.  You can continue to update v0 and v1
-after you have pushed v3.
+Example of use:
 
-Markdown is supported, following the CommonMark specification.
+```python
+class UbuntuCharm(CharmBase):
+    def __init__(self, *args):
+        ...
+        self.framework.observe(self.on.install, self._on_install)
+        self.framework.observe(self.on.update_status, self._on_update_status)
+        self.framework.observe(self.on.remove, self._on_remove)
+        self.grub = grub.GrubConfig(self.meta.name)
+        log.debug("found keys %s in GRUB config file", self.grub.keys())
+
+    def _on_install(self, _):
+        try:
+            self.grub.update(
+                {"GRUB_CMDLINE_LINUX_DEFAULT": "$GRUB_CMDLINE_LINUX_DEFAULT hugepagesz=1G"}
+            )
+        except grub.ValidationError as error:
+            self.unit.status = BlockedStatus(f"[{error.key}] {error.message}")
+
+    def _on_update_status(self, _):
+        if self.grub["GRUB_CMDLINE_LINUX_DEFAULT"] != "$GRUB_CMDLINE_LINUX_DEFAULT hugepagesz=1G":
+            self.unit.status = BlockedStatus("wrong GRUB configuration")
+
+    def _on_remove(self, _):
+        self.grub.remove()
+```
 """
 
+import filecmp
+import io
+import logging
+import os
+import shlex
+import subprocess
+from pathlib import Path
+from typing import Dict, Mapping, Optional, Set, Tuple
+
+logger = logging.getLogger(__name__)
+
 # The unique Charmhub library identifier, never change it
 LIBID = "1f73a0e0c78349bc88850022e02b33c7"
 
@@ -29,4 +72,311 @@
 # to 0 if you are raising the major API version
 LIBPATCH = 1
 
-# TODO: add your code here! Happy coding!
+GRUB_DIRECTORY = Path("/etc/default/grub.d/")
+CHARM_CONFIG_PREFIX = "90-juju"
+GRUB_CONFIG = GRUB_DIRECTORY / "95-juju-charm.cfg"
+CONFIG_HEADER = f"""# This config file was produced by GRUB lib v{LIBAPI}.{LIBPATCH}.
+# https://charmhub.io/operator-libs-linux/libraries/grub
+"""
+FILE_LINE_IN_DESCRIPTION = "#   {path}"
+CONFIG_DESCRIPTION = """
+# This file represents the output of the GRUB lib, which can combine multiple
+# configurations into a single file like this.
+#
+# Original files:
+{configs}
+#
+# If you change this file, run 'update-grub' afterwards to update
+# /boot/grub/grub.cfg.
+# For full documentation of the options in this file, see:
+#   info -f grub -n 'Simple configuration'
+"""
+
+
+class ValidationError(ValueError):
+    """Exception representing value validation error."""
+
+    def __init__(self, key: str, message: str) -> None:
+        super().__init__(message)
+        self.key = key
+        self.message = message
+
+    def __str__(self) -> str:
+        """Return string representation of error."""
+        return self.message
+
+
+class IsContainerError(Exception):
+    """Exception if local machine is container."""
+
+
+class ApplyError(Exception):
+    """Exception if applying new config failed."""
+
+
+def _split_config_line(line: str) -> Tuple[str, str]:
+    """Split GRUB config line to obtain key and value."""
+    key, raw_value = line.split("=", 1)
+    value, *not_expected_values = shlex.split(raw_value)
+    if not_expected_values:
+        logger.error("unexpected value %s for %s key", raw_value, key)
+        raise ValueError(f"unexpected value {raw_value} for {key} key")
+
+    return key, value
+
+
+def _parse_config(stream: io.TextIOWrapper) -> Dict[str, str]:
+    """Parse config file lines."""
+    config = {}
+    for line in stream:
+        line = line.strip()
+        if not line or line.startswith("#"):
+            logger.debug("skipping line `%s`", line)
+            continue
+
+        key, value = _split_config_line(line)
+        if key in config:
+            logger.warning("key %s is duplicated in config", key)
+
+        config[key] = value
+
+    return config
+
+
+def _load_config(path: Path) -> Dict[str, str]:
+    """Load config file from /etc/default/grub.d/ directory."""
+    if not path.exists():
+        raise FileNotFoundError("GRUB config file %s was not found", path)
+
+    with open(path, "r", encoding="UTF-8") as file:
+        config = _parse_config(file)
+
+    logger.info("GRUB config file %s was loaded", path)
+    logger.debug("config file %s", config)
+    return config
+
+
+def _save_config(path: Path, config: Dict[str, str], header: str = CONFIG_HEADER) -> None:
+    """Save GRUB config file."""
+    if path.exists():
+        logger.debug("GRUB config %s already exist and it will overwritten", path)
+
+    context = [f"{key}={shlex.quote(value)}" for key, value in config.items()]
+    with open(path, "w", encoding="UTF-8") as file:
+        file.writelines([header, *context])
+
+    logger.info("GRUB config file %s was saved", path)
+
+
+def check_update_grub() -> bool:
+    """Report whether an update to /boot/grub/grub.cfg is available."""
+    main_grub_cfg = Path("/boot/grub/grub.cfg")
+    tmp_path = Path("/tmp/tmp_grub.cfg")
+    try:
+        subprocess.check_call(
+            ["/usr/sbin/grub-mkconfig", "-o", f"{tmp_path}"], stderr=subprocess.STDOUT
+        )
+    except subprocess.CalledProcessError as error:
+        logger.exception(error)
+        raise
+
+    return not filecmp.cmp(main_grub_cfg, tmp_path)
+
+
+def is_container() -> bool:
+    """Report whether the local machine is a container."""
+    try:
+        output = subprocess.check_output(
+            ["/usr/bin/systemd-detect-virt", "--container"], stderr=subprocess.STDOUT
+        ).decode()
+        logger.debug("detect virt type %s", output)
+        return True
+    except subprocess.CalledProcessError:
+        return False
+
+
+class Config(Mapping[str, str]):
+    """Manages GRUB configuration.
+
+    This object will load current configuration option for GRUB and provide option
+    to update it with simple validation, remove charm option and apply those changes.
+    """
+
+    _lazy_data: Optional[Dict[str, str]] = None
+
+    def __init__(self, charm_name: str) -> None:
+        """Initialize the GRUB config."""
+        self._charm_name = charm_name
+
+    def __contains__(self, key: str) -> bool:
+        """Check if key is in config."""
+        return key in self._data
+
+    def __len__(self):
+        """Get size of config."""
+        return len(self._data)
+
+    def __iter__(self):
+        """Iterate over config."""
+        return iter(self._data)
+
+    def __getitem__(self, key: str) -> str:
+        """Get value for key form config."""
+        return self._data[key]
+
+    @property
+    def _data(self) -> Dict[str, str]:
+        """Data property."""
+        if self._lazy_data is None:
+            try:
+                self._lazy_data = _load_config(GRUB_CONFIG)
+            except FileNotFoundError:
+                logger.debug("there is no GRUB config file %s yet", GRUB_CONFIG)
+                self._lazy_data = {}
+
+        return self._lazy_data
+
+    def _save_grub_configuration(self) -> None:
+        """Save current GRUB configuration."""
+        logger.info("saving new GRUB config to %s", GRUB_CONFIG)
+        applied_configs = {self.path, *self.applied_configs}  # using set to drop duplicity
+        registered_configs = os.linesep.join(
+            FILE_LINE_IN_DESCRIPTION.format(path=path) for path in applied_configs
+        )
+        header = CONFIG_HEADER + CONFIG_DESCRIPTION.format(configs=registered_configs)
+        _save_config(GRUB_CONFIG, self._data, header)
+
+    def _set_value(self, key: str, value: str, blocked_keys: Set[str]) -> bool:
+        """Set new value for key."""
+        logger.debug("[%s] setting new value %s for key %s", self.charm_name, value, key)
+        current_value = self._data.get(key)
+        if current_value == value:
+            return False
+
+        # validation
+        if key in self and current_value != value and key in blocked_keys:
+            logger.error(
+                "[%s] tries to overwrite key %s, which has value %s, with value %s",
+                self.charm_name,
+                key,
+                current_value,
+                value,
+            )
+            raise ValidationError(
+                key, f"key {key} already exists and its value is {current_value}"
+            )
+
+        self._data[key] = value
+        return True
+
+    def _update(self, config: Dict[str, str]) -> Set[str]:
+        """Update data in object."""
+        logger.debug("[%s] updating current config", self.charm_name)
+        changed_keys = set()
+        blocked_keys = self.blocked_keys
+        for key, value in config.items():
+            changed = self._set_value(key, value, blocked_keys)
+            if changed:
+                changed_keys.add(key)
+
+        return changed_keys
+
+    @property
+    def applied_configs(self) -> Dict[Path, Dict[str, str]]:
+        """Return list of charms configs which registered config in LIB_CONFIG_DIRECTORY."""
+        configs = {}
+        for path in sorted(GRUB_DIRECTORY.glob(f"{CHARM_CONFIG_PREFIX}-*")):
+            configs[path] = _load_config(path)
+            logger.debug("load config file %s", path)
+
+        return configs
+
+    @property
+    def blocked_keys(self) -> Set[str]:
+        """Get set of configured keys by other charms."""
+        return {
+            key
+            for path, config in self.applied_configs.items()
+            if path != self.path
+            for key in config
+        }
+
+    @property
+    def charm_name(self) -> str:
+        """Get charm name or use value obtained from JUJU_UNIT_NAME env."""
+        return self._charm_name
+
+    @property
+    def path(self) -> Path:
+        """Return path for charm config."""
+        return GRUB_DIRECTORY / f"{CHARM_CONFIG_PREFIX}-{self.charm_name}"
+
+    def apply(self):
+        """Check if an update to /boot/grub/grub.cfg is available."""
+        if not check_update_grub():
+            logger.info("[%s] no available GRUB updates found", self.charm_name)
+            return
+
+        try:
+            subprocess.check_call(["/usr/sbin/update-grub"], stderr=subprocess.STDOUT)
+        except subprocess.CalledProcessError as error:
+            logger.error(
+                "[%s] applying GRUB config failed with errors: %s", self.charm_name, error.stdout
+            )
+            raise ApplyError("New config check failed.") from error
+
+    def remove(self, apply: bool = True) -> Set[str]:
+        """Remove config for charm.
+
+        This function will remove config file for charm and re-create the `95-juju-charm.cfg`
+        GRUB config file without changes made by this charm.
+        """
+        if not self.path.exists():
+            logger.debug("[%s] there is no charm config file %s", self.charm_name, self.path)
+            return set()
+
+        self.path.unlink()
+        logger.info("[%s] charm config file %s was removed", self.charm_name, self.path)
+        config = {}
+        for _config in self.applied_configs.values():
+            config.update(_config)
+
+        changed_keys = set(self._data) - set(config.keys())
+        self._lazy_data = config
+        self._save_grub_configuration()
+        if apply:
+            self.apply()
+
+        return changed_keys
+
+    def update(self, config: Dict[str, str], apply: bool = True) -> Set[str]:
+        """Update the Grub configuration."""
+        if is_container():
+            raise IsContainerError("Could not configure GRUB config on container.")
+
+        snapshot = self._data.copy()
+        try:
+            changed_keys = self._update(config)
+            if changed_keys:
+                self._save_grub_configuration()
+                if apply:
+                    self.apply()
+        except ValidationError as error:
+            logger.error("[%s] validation failed with message: %s", self.charm_name, error.message)
+            self._lazy_data = snapshot
+            logger.info("[%s] restored snapshot for Config object", self.charm_name)
+            raise
+        except ApplyError as error:
+            logger.error(
+                "[%s] applying new GRUB config failed with error: %s", self.charm_name, error
+            )
+            self._lazy_data = snapshot
+            self._save_grub_configuration()  # save snapshot copy of grub config
+            logger.info(
+                "[%s] restored snapshot for Config object and GRUB configuration", self.charm_name
+            )
+            raise
+
+        logger.debug("[%s] saving copy of charm config to %s", self.charm_name, GRUB_DIRECTORY)
+        _save_config(self.path, config)
+        return changed_keys