#!/usr/bin/python3
#
# SPDX-License-Identifier: GPL-2.0-only
# (c) 2026 Gerd Hoffmann
#
# pylint: disable=too-many-arguments
# pylint: disable=too-many-positional-arguments
""" generate cloud init iso images """
import os
import sys
import shutil
import pathlib
import tempfile
import subprocess

import yaml

try:
    import crypt_r as crypt
except ImportError:
    import crypt


class LiteralStr(str):
    """Marker class: dump as literal block scalar (|)."""

def literal_representer(dumper, data):
    return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")

class CloudInitDumper(yaml.SafeDumper):
    """Safe YAML dumper with support for LiteralStr literal block scalars."""

CloudInitDumper.add_representer(LiteralStr, literal_representer)


class CloudInitISO:
    """ cloud init iso image generator """

    def __init__(self, hostname):
        self.config = {
            "hostname": hostname,
        }
        self.meta = {
            "local-hostname": hostname
        }

    @staticmethod
    def crypt(passwd):
        return crypt.crypt(passwd, crypt.mksalt(method=crypt.METHOD_SHA512))

    def add_user(self, name, passwd = None, ssh_keys = None, sudo = False, shell = "/bin/bash"):
        user = {
            "name"  : name,
            "shell" : shell
        }
        if passwd:
            user["passwd"] = self.crypt(passwd)
            user["lock_passwd"] = False
        if ssh_keys:
            user["ssh_authorized_keys"] = ssh_keys
        if sudo:
            user["sudo"] = ["ALL=(ALL) NOPASSWD:ALL"]
        if not self.config.get("users"):
            self.config["users"] = []
        self.config["users"].append(user)

    def add_package(self, package_name):
        """ package to install """
        if not self.config.get("packages"):
            self.config["packages"] = []
        self.config["packages"].append(package_name)

    def add_runcmd(self, command):
        """ command to run on first boot """
        if not self.config.get("runcmd"):
            self.config["runcmd"] = []
        self.config["runcmd"].append(command)

    def add_bootcmd(self, command):
        """ command to run on every boot """
        if not self.config.get("bootcmd"):
            self.config["bootcmd"] = []
        self.config["bootcmd"].append(command)

    def add_file_content(self, filename, content, owner = 'root:root'):
        """ file to write to disk """
        if not self.config.get("write_files"):
            self.config["write_files"] = []
        wf = {
            "path"    : filename,
            "owner"   : owner,
            "content" : content,
        }
        self.config["write_files"].append(wf)

    def add_file_copy(self, filename):
        try:
            with open(filename, 'r', encoding = 'utf-8') as f:
                content = f.read()
            self.add_file_content(filename, LiteralStr(content))
        except UnicodeDecodeError:
            with open(filename, 'rb') as f:
                content = f.read()
            self.add_file_content(filename, content)

    # pylint: disable=consider-using-with
    def create_iso(self):
        iso = tempfile.NamedTemporaryFile(suffix = '.iso')
        with tempfile.TemporaryDirectory() as tmpdir:
            tmpdir = pathlib.Path(tmpdir)

            user_data_file = tmpdir / "user-data"
            with user_data_file.open("w") as f:
                f.write("#cloud-config\n")
                yaml.dump(self.config, f, default_flow_style=False,
                          sort_keys=False, Dumper=CloudInitDumper)

            meta_data_file = tmpdir / "meta-data"
            with meta_data_file.open("w") as f:
                yaml.dump(self.meta, f, default_flow_style=False,
                          sort_keys=False, Dumper=CloudInitDumper)

            tmtdata = os.environ.get('TMT_TEST_DATA')
            if tmtdata:
                dstdir = pathlib.Path(tmtdata) / "cloud-init"
                os.makedirs(dstdir, exist_ok = True)
                shutil.copy(user_data_file, dstdir / "user-data")
                shutil.copy(meta_data_file, dstdir / "meta-data")

            subprocess.run([
                "xorrisofs",
                "-output", iso.name,
                "-volid", "cidata",
                "-joliet",
                "-rock",
                str(user_data_file),
                str(meta_data_file)
            ], capture_output = True, check=True)

        return iso

    def dump_user(self):
        y = yaml.dump(self.config, default_flow_style=False,
                      sort_keys=False, Dumper=CloudInitDumper)
        print(y, end='')

def main():
    c = CloudInitISO("test-host")
    c.add_user("user", "passwd", sudo = True)
    c.add_package("pciutils")
    c.add_bootcmd("echo hello bootcmd")
    c.add_runcmd("echo hello runcmd #1")
    c.add_runcmd("echo hello runcmd #2")
    c.add_file_content("/root/test", "hello world")
    c.add_file_copy("/etc/resolv.conf")
    c.dump_user()

if __name__ == '__main__':
    sys.exit(main())
