From feeae04835574655aff2397dd78a5e1522c5156b Mon Sep 17 00:00:00 2001
From: Lili Nie <lnie@redhat.com>
Date: Sun, 14 Jul 2024 21:49:05 -0400
Subject: [PATCH] feat: Add support for translanting %pre and main body part of
 kickstart

---
 src/mrack/transformers/beaker.py      | 37 +++++++++++++++++++++------
 tests/unit/data/key_one               |  1 +
 tests/unit/data/key_two               |  1 +
 tests/unit/test_beaker_transformer.py | 32 +++++++++++++++++++++--
 4 files changed, 61 insertions(+), 10 deletions(-)
 create mode 100644 tests/unit/data/key_one
 create mode 100644 tests/unit/data/key_two

diff --git a/src/mrack/transformers/beaker.py b/src/mrack/transformers/beaker.py
index 6a2816af..93776376 100644
--- a/src/mrack/transformers/beaker.py
+++ b/src/mrack/transformers/beaker.py
@@ -86,17 +86,37 @@ def _get_ks_meta(self, host):
 
     def _construct_ks_append_script(self, ks_append, pubkeys=None):
         """Create ks_appdend from requirements."""
+        res_ks_list = []
         if not ks_append and not pubkeys:
             return []
 
-        res_ks_append = ["%post"]
-        if pubkeys:
-            ks_append += self._allow_ssh_keys(pubkeys)
+        if isinstance(ks_append, dict):
+            res_ks_pre = ks_append.get("pre-install")
+            res_ks = ks_append.get("script")
+            res_ks_post = ks_append.get("post-install")
+            if res_ks_pre:
+                if res_ks_pre.startswith("%pre"):
+                    res_ks_list += [res_ks_pre]
+                else:
+                    res_ks_list += ["%pre"] + [res_ks_pre] + ["%end"]
+            if res_ks:
+                res_ks_list += [res_ks]
+            if res_ks_post:
+                if res_ks_post.startswith("%post"):
+                    res_ks_list += [res_ks_post]
+                else:
+                    res_ks_list += ["%post"] + [res_ks_post] + ["%end"]
+        else:
+            res_ks_list = ["%post"]
+            res_ks_list += ks_append
+            res_ks_list.append("%end")
 
-        res_ks_append += ks_append
+        if pubkeys:
+            res_ks_list += ["%post"] + self._allow_ssh_keys(pubkeys) + ["%end"]
 
-        res_ks_append.append("%end")
-        return ["\n".join(res_ks_append)]
+        return [
+            "\n".join(res_ks_list),
+        ]
 
     def _allow_ssh_keys(self, pubkeys):
         """Create ssh key content to be injected to xml."""
@@ -104,8 +124,9 @@ def _allow_ssh_keys(self, pubkeys):
         keys_content.append("mkdir -p /root/.ssh")
         keys_content.append('cat >>/root/.ssh/authorized_keys << "__EOF__"')
         keys_content.append("# keys added by mrack:")
-
-        for key in set(pubkeys):
+        keys = list(set(pubkeys))
+        keys.sort(key=pubkeys.index)
+        for key in keys:
             with open(os.path.expanduser(key), "r", encoding="utf-8") as key_file:
                 keys_content.append(f"{key_file.read().strip()}")
 
diff --git a/tests/unit/data/key_one b/tests/unit/data/key_one
new file mode 100644
index 00000000..6ff47f42
--- /dev/null
+++ b/tests/unit/data/key_one
@@ -0,0 +1 @@
+ssh-rsa key_one_content
diff --git a/tests/unit/data/key_two b/tests/unit/data/key_two
new file mode 100644
index 00000000..2e21ab90
--- /dev/null
+++ b/tests/unit/data/key_two
@@ -0,0 +1 @@
+ssh-rsa key_two_content
diff --git a/tests/unit/test_beaker_transformer.py b/tests/unit/test_beaker_transformer.py
index 54a88366..252289f3 100644
--- a/tests/unit/test_beaker_transformer.py
+++ b/tests/unit/test_beaker_transformer.py
@@ -5,6 +5,7 @@
 from mrack.providers.beaker import BeakerProvider
 
 from .mock_data import MockedBeakerTransformer, provisioning_config
+from .utils import get_file_path
 
 
 class TestBeakerTransformer:
@@ -27,7 +28,28 @@ class TestBeakerTransformer:
     ]
     default_retention_tag = "audit"
     default_product = "[internal]"
-
+    key_one = get_file_path("key_one")
+    key_two = get_file_path("key_two")
+    raw_maximal_ksappend = (
+        [
+            "%pre\npre_dummy\n%end\nscript_dummy\n%post\npost_dummy\n%end",
+        ]
+        + [
+            '%post\nmkdir -p /root/.ssh\ncat >>/root/.ssh/authorized_keys << "__EOF__"',
+        ]
+        + [
+            "# keys added by mrack:\nssh-rsa key_one_content\nssh-rsa key_two_content",
+        ]
+        + [
+            "# end section of keys added by mrack\n__EOF__\nrestorecon -R /root/.ssh",
+        ]
+        + [
+            "chmod go-w /root /root/.ssh /root/.ssh/authorized_keys\n%end",
+        ]
+    )
+    maximal_ksappend = [
+        "\n".join(raw_maximal_ksappend),
+    ]
     fedora = {
         "name": f"fedora.{domain_name}",
         "role": "client",
@@ -36,6 +58,12 @@ class TestBeakerTransformer:
         "restraint_id": 1,
         "beaker": {
             "ks_meta": "FEDORA_HOST_KS_META",
+            "ks_append": {
+                "pre-install": "%pre\npre_dummy\n%end",
+                "script": "script_dummy",
+                "post-install": "%post\npost_dummy\n%end",
+            },
+            "pubkeys": [key_one, key_two],
             "tasks": [
                 {
                     "name": "/distribution/check-install",
@@ -164,7 +192,7 @@ async def create_transformer(self, legacy=False):
                     "distro": "Fedora-36%",
                     "variant": "Server",
                     "ks_meta": "FEDORA_HOST_KS_META",
-                    "ks_append": default_ks_append,
+                    "ks_append": maximal_ksappend,
                     "whiteboard": default_whiteboard,
                     "priority": default_prio,
                     "tasks": [