ansible-systemd/plugins/module_utils/module.py

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

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 mo"""

    name: ClassVar[str]
    module: basic.AnsibleModule
    msg: str
    result: dict
    module_spec: ClassVar[dict]
    _common_args = dict()

    @property
    def params(self) -> Dict[str, Any]:
        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.msg = ""
        self.tmpdir = pathlib.Path(self.module.tmpdir)

    def set(self, key: str, value):
        self.result[key] = value

    def diff(
        self,
        before,
        after,
        before_header: Optional[str] = None,
        after_header: Optional[str] = None,
    ):
        if "diff" not in self.result:
            self.result["diff"] = list()
        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
        self.result["diff"].append(diff)

    def get(self, key: str, default: T = None) -> T:
        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):
        return self.result["changed"]

    @changed.setter
    def changed_set(self, value):
        self.result["changed"] = not not value

    def prepare(self):
        raise NotImplementedError()

    def check(self):
        raise NotImplementedError()

    def run(self):
        raise NotImplementedError()

    def __call__(self):
        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:
        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):
    unitfile: pathlib.Path
    _common_args = dict(
        supports_check_mode=True,
        add_file_common_args=True,
    )
    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():
            if "diff" not in self.result:
                self.result["diff"] = list()
            diff = dict()
            self.result["changed"] = self.module.set_owner_if_different(
                self.unitfile.as_posix(),
                "root",
                self.result["changed"],
                diff,
            )
            self.result["changed"] = self.module.set_group_if_different(
                self.unitfile.as_posix(),
                "root",
                self.result["changed"],
                diff,
            )
            self.result["changed"] = self.module.set_mode_if_different(
                self.unitfile.as_posix(),
                "0644",
                self.result["changed"],
                diff,
            )
            self.result["diff"].append(diff)

    def check(self):
        if "changed" in self.result:
            changed = self.result["changed"]
        else:
            changed = False
        self.unitfile_gen()
        if not self.unitfile.exists():
            self.diff("", self.unit(), self.unitfile.as_posix())
            changed = True
        else:
            if self.module.sha256(self.unitfile.as_posix()) != self.module.sha256(
                (self.tmpdir / "newunit").as_posix()
            ):
                changed = True
                self.diff(
                    before=self.unitfile.read_text(),
                    after=self.unit(),
                    before_header=self.unitfile.as_posix(),
                )
        self.set("changed", changed)
        if hasattr(self, "post") and self.post is not None:
            self.post()
        return changed

    def run(self):
        if not self.check():
            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()