From 139ab4f67a657e366d58fabfcaeb841e61b0d26e Mon Sep 17 00:00:00 2001
From: Aditya Nagesh <adityanagesh@microsoft.com>
Date: Sat, 9 Nov 2024 16:47:49 +0530
Subject: [PATCH] Kernel Installer: Add B4 Modifier

---
 lisa/tools/__init__.py                       |  2 ++
 lisa/tools/b4.py                             | 29 ++++++++++++++++++++
 lisa/tools/git.py                            | 11 ++++++++
 lisa/transformers/kernel_source_installer.py | 29 +++++++++++++++++++-
 4 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 lisa/tools/b4.py

diff --git a/lisa/tools/__init__.py b/lisa/tools/__init__.py
index 720f3709f4..4bef1fd595 100644
--- a/lisa/tools/__init__.py
+++ b/lisa/tools/__init__.py
@@ -16,6 +16,7 @@
 )
 
 from .aria import Aria
+from .b4 import B4
 from .blkid import Blkid
 from .bzip2 import Bzip2
 from .cargo import Cargo
@@ -128,6 +129,7 @@
 __all__ = [
     "AptAddRepository",
     "Aria",
+    "B4",
     "Blkid",
     "Bzip2",
     "Cargo",
diff --git a/lisa/tools/b4.py b/lisa/tools/b4.py
new file mode 100644
index 0000000000..a97a6c8f73
--- /dev/null
+++ b/lisa/tools/b4.py
@@ -0,0 +1,29 @@
+import pathlib
+import re
+from lisa.executable import Tool
+from lisa.operating_system import Posix
+from lisa.util import get_matched_str
+
+
+class B4(Tool):
+    _output_file_pattern = re.compile(r"[\w-]+\.mbx")
+
+    @property
+    def command(self) -> str:
+        return "b4"
+
+    @property
+    def can_install(self) -> bool:
+        return True
+
+    def _install(self) -> bool:
+        if isinstance(self.node.os, Posix):
+            self.node.os.install_packages("b4")
+        return self._check_exists()
+
+    def am(self, message_id: str, output_dir: pathlib.PurePath) -> pathlib.PurePath:
+        result = self.run(
+            f"am -o '{output_dir}' '{message_id}'", force_run=True, expected_exit_code=0
+        )
+        path_str = get_matched_str(result.stdout, self._output_file_pattern)
+        return pathlib.PurePath(output_dir, path_str)
diff --git a/lisa/tools/git.py b/lisa/tools/git.py
index d1577f404e..d55def77d3 100644
--- a/lisa/tools/git.py
+++ b/lisa/tools/git.py
@@ -201,6 +201,17 @@ def apply(
         )
         result.assert_exit_code(message=f"failed on applying patches. {result.stdout}")
 
+    def am(self, cwd: pathlib.PurePath, patch: pathlib.PurePath) -> None:
+        result = self.run(
+            f"am {patch}",
+            shell=True,
+            cwd=cwd,
+            force_run=True,
+            no_info_log=True,
+            no_error_log=True,
+        )
+        result.assert_exit_code(message=f"failed on applying patches. {result.stdout}")
+
     def list_tags(self, cwd: pathlib.PurePath) -> List[str]:
         result = self.run(
             "--no-pager tag --color=never",
diff --git a/lisa/transformers/kernel_source_installer.py b/lisa/transformers/kernel_source_installer.py
index 577e3975e3..bf11af927c 100644
--- a/lisa/transformers/kernel_source_installer.py
+++ b/lisa/transformers/kernel_source_installer.py
@@ -11,7 +11,7 @@
 from lisa.base_tools import Mv
 from lisa.node import Node
 from lisa.operating_system import CBLMariner, Redhat, Ubuntu
-from lisa.tools import Cp, Echo, Git, Make, Sed, Uname
+from lisa.tools import Cp, Echo, Git, Make, Sed, Uname, B4
 from lisa.tools.gcc import Gcc
 from lisa.tools.lscpu import Lscpu
 from lisa.util import LisaException, field_metadata, subclasses
@@ -75,6 +75,12 @@ class PatchModifierSchema(BaseModifierSchema):
     file_pattern: str = "*.patch"
 
 
+@dataclass_json()
+@dataclass
+class b4PatchModifierSchema(BaseModifierSchema):
+    message_id: str = field(default="", metadata=field_metadata(required=True))
+
+
 @dataclass_json()
 @dataclass
 class SourceInstallerSchema(BaseInstallerSchema):
@@ -489,3 +495,24 @@ def _get_code_path(path: str, node: Node, default_name: str) -> PurePath:
         code_path = node.working_path / default_name
 
     return code_path
+
+
+class b4PatchModifier(BaseModifier):
+    @classmethod
+    def type_name(cls) -> str:
+        return "b4_patch"
+
+    @classmethod
+    def type_schema(cls) -> Type[schema.TypedSchema]:
+        return b4PatchModifierSchema
+
+    def modify(self) -> None:
+        runbook: b4PatchModifierSchema = self.runbook
+
+        git = self._node.tools[Git]
+        b4 = self._node.tools[B4]
+
+        message_id = runbook.message_id
+        patch_file = b4.am(message_id=message_id, output_dir=self._code_path)
+
+        git.am(cwd=self._code_path, patch=patch_file)