From e1e3d0a7d40256854f468c7a8d36bde7e9f39974 Mon Sep 17 00:00:00 2001
From: Lili Nie <lnie@redhat.com>
Date: Mon, 1 Jul 2024 02:53:27 -0400
Subject: [PATCH] feat: Add support for translanting %pre and main body part of
 kickstart

---
 src/mrack/transformers/beaker.py      | 30 +++++++++++++++++++++------
 tests/unit/test_beaker_transformer.py | 22 +++++++++++++++++++-
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/src/mrack/transformers/beaker.py b/src/mrack/transformers/beaker.py
index 1f504c00..d2edcb8f 100644
--- a/src/mrack/transformers/beaker.py
+++ b/src/mrack/transformers/beaker.py
@@ -86,17 +86,35 @@ 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."""
diff --git a/tests/unit/test_beaker_transformer.py b/tests/unit/test_beaker_transformer.py
index bbaf0a69..1a19ff26 100644
--- a/tests/unit/test_beaker_transformer.py
+++ b/tests/unit/test_beaker_transformer.py
@@ -19,6 +19,20 @@ class TestBeakerTransformer:
     default_tasks = [{"name": "/distribution/dummy", "role": "STANDALONE"}]
     default_retention_tag = "audit"
     default_product = "[internal]"
+    maximal_ksappend = """
+    %pre\npre_dummy\n%end\nscript_dummy\n%post\npost_dummy\n%end
+    %post
+    mkdir -p /root/.ssh
+    cat >>/root/.ssh/authorized_keys << "__EOF__"
+    # keys added by mrack:
+    key_one
+    key_two
+    # end section of keys added by mrack
+    __EOF__
+    restorecon -R /root/.ssh
+    chmod go-w /root /root/.ssh /root/.ssh/authorized_keys
+    %end
+    """
 
     fedora = {
         "name": f"fedora.{domain_name}",
@@ -28,6 +42,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",
@@ -156,7 +176,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": [