ansible-systemd/plugins/module_utils/module.py

import pathlib
from typing import Any, Callable, ClassVar, Dict, Optional, TypeVar, NoReturn, overload

import ansible.module_utils.basic as basic
from ansible.module_utils.generic import _sdict

__all__ = (
    "AnsibleModule",
    "SystemdUnitModule",
)


T = TypeVar("T")


class AnsibleModule(object):
    """Simple wrapper for the basic.AnsibleModule"""

    #: name of the module. This is required for the generation of the Ansible documentation
    name: ClassVar[str]
    #: The AnsibleModule for this Module
    module: basic.AnsibleModule
    #: The result of this Module call. It always contains the changed key, so in any case an Module can report if it changed anything
    result: dict
    #: the specification of the arguments. Subclasses that are usable Modules must set this value.
    module_spec: ClassVar[dict]
    #: This is set by classes that define common things for their subclasses, like behaviour of the run and check methods. This is used by `SystemdUnitModule`
    _common_args = dict()

    @property
    def params(self) -> Dict[str, Any]:
        """params is an wrapper for the module.params"""
        return self.module.params  # type: ignore

    def __init__(self):
        self.result = dict(changed=False)
        specs = dict()
        specs.update(self._common_args)
        specs.update(self.module_spec)
        self.module = basic.AnsibleModule(**specs)
        self.tmpdir = pathlib.Path(self.module.tmpdir)

    def set(self, key: str, value):
        """sets an value for the result"""
        self.result[key] = value

    @overload
    def diff(self, diff: Dict[str, str]):
        pass

    def diff(
        self,diff: Optional[Dict[str, str]] = None,*,
        before: Optional[str] = None,
        after: Optional[str] = None,
        before_header: Optional[str] = None,
        after_header: Optional[str] = None,
    ):
        """adds the special return value "diff". This allows Modules to present the changes of files to the caller. it takes care of the special semantics of the return value"""
        if "diff" not in self.result:
            self.result["diff"] = list()
        if diff is not None and not any((before is not None, after is not None)):
            pass
        elif all((before is not None, after is not None, diff is None)):
            diff = dict(
                before=before,
                after=after,
            )
            if before_header is not None:
                diff["before_header"] = before_header
            if after_header is not None:
                diff["after_header"] = after_header
        else:
            raise TypeError("only diff or before and after can be set, not both of them")
        self.result["diff"].append(diff)

    def get(self, key: str, default: T = None) -> T:
        """returns an Parameter of the Module."""
        if key not in self.params.keys():
            return default
        if self.params[key] is None and default is not None:
            return default
        if self.params[key] is None or key not in self.params:
            raise KeyError()
        return self.params[key]

    @property
    def changed(self):
        """returns if changes were detected/made"""
        return self.result["changed"]

    @changed.setter
    def changed_set(self, value):
        """sets the changed value. this is always converted to bool"""
        self.result["changed"] = not not value

    def prepare(self):
        raise NotImplementedError()

    def check(self):
        raise NotImplementedError()

    def run(self):
        raise NotImplementedError()

    def __call__(self) -> NoReturn:
        """This calls the module. first prepare is called and then check or run, depending on the check mode.
        If an exception is raised this is catched and the module automatically fails with an traceback"""
        self.prepare()
        try:
            if self.module.check_mode:
                self.check()
            else:
                self.run()
        except Exception as exc:
            import traceback

            self.module.fail_json(
                "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
                **self.result,
            )
        self.module.exit_json(**self.result)

    @classmethod
    def doc(cls) -> str:
        """this returns the documentation string of the module. If the help arguments of an Types method was given, it adds this as an helptext of this parameter"""
        try:
            import yaml
        except ImportError:
            return "---\n"
        doc = cls.__doc__
        if doc is None:
            doc = ""
        options = dict()
        help: _sdict
        for option, help in cls.module_spec["argument_spec"].items():
            options[option] = dict(
                type=help["type"],
            )
            if hasattr(help, "_help") and help._help is not None:
                options[option]["description"] = help._help.split("\n")
            if "required" in help and help["required"]:
                options[option]["required"] = True
            else:
                options[option]["required"] = False
            if help["type"] == "list":
                options[option]["elements"] = help["elements"]
                if not options[option]["required"]:
                    options[option]["default"] = []
            if "default" in help:
                options[option]["default"] = help["default"]
            if "choices" in help:
                options[option]["choices"] = tuple(help["choices"])
        docu = doc.split("\n")
        return str(
            yaml.safe_dump(
                dict(
                    module=cls.name,
                    short_description=docu[0],
                    description=docu,
                    options=options,
                ),
                stream=None,
                explicit_start="---",
            )
        )


class SystemdUnitModule(AnsibleModule):
    #: path of the unitfile managed by this module
    unitfile: pathlib.Path
    #: subclasses of this always support the file common args and the check mode
    _common_args = dict(
        supports_check_mode=True,
        add_file_common_args=True,
    )
    #: if defined it will be called after run has changed the unitfile
    post: Optional[Callable[[], None]]

    def unit(self) -> str:
        raise NotImplementedError()

    def unitfile_gen(self):
        path = self.tmpdir / "newunit"
        with open(path, "w") as unit:
            unit.write(self.unit())
        self.module.set_owner_if_different(path.as_posix(), "root", False)
        self.module.set_group_if_different(path.as_posix(), "root", False)
        self.module.set_mode_if_different(path.as_posix(), "0644", False)
        if self.unitfile.exists():
            diff = dict()
            self.changed = self.changed | self.module.set_owner_if_different(
                self.unitfile.as_posix(),
                "root",
                self.result["changed"],
                diff,
            )
            self.diff(diff)
            diff = dict()
            self.changed = self.changed | self.module.set_group_if_different(
                self.unitfile.as_posix(),
                "root",
                self.result["changed"],
                diff,
            )
            self.diff(diff)
            diff = dict()
            self.changed = self.changed | self.module.set_mode_if_different(
                self.unitfile.as_posix(),
                "0644",
                self.result["changed"],
                diff,
            )
            self.diff(diff)

    def check(self):
        self.unitfile_gen()
        if not self.unitfile.exists():
            self.diff(before="", after=self.unit(), before_header=self.unitfile.as_posix())
            self.changed = True
        else:
            if self.module.sha256(self.unitfile.as_posix()) != self.module.sha256(
                (self.tmpdir / "newunit").as_posix()
            ):
                self.changed = True
                self.diff(
                    before=self.unitfile.read_text(),
                    after=self.unit(),
                    before_header=self.unitfile.as_posix(),
                )

    def run(self):
        self.check()
        if not self.changed:
            return
        self.module.atomic_move(
            src=(self.tmpdir / "newunit").as_posix(),
            dest=self.unitfile.as_posix(),
        )
        if hasattr(self, "post") and self.post is not None:
            self.post()