From 8baa196a8d2e8e17f23195b74f379de70b15961c Mon Sep 17 00:00:00 2001
From: Aditya Nagesh <adityanagesh@microsoft.com>
Date: Tue, 12 Nov 2024 15:15:51 +0530
Subject: [PATCH] Add B4 Tool

---
 lisa/tools/__init__.py |  2 ++
 lisa/tools/b4.py       | 47 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 lisa/tools/b4.py

diff --git a/lisa/tools/__init__.py b/lisa/tools/__init__.py
index 0516a6cae6..b9422e96bf 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
@@ -129,6 +130,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..5f364bfd8a
--- /dev/null
+++ b/lisa/tools/b4.py
@@ -0,0 +1,47 @@
+import pathlib
+import re
+from lisa.executable import Tool
+from lisa.operating_system import Debian
+from lisa.tools.python import Pip
+from lisa.util import LisaException, find_group_in_lines
+
+
+class B4(Tool):
+    # Output log is of the form
+    # git am /mnt/code/linux/v2_20241029_xx_drivers_xx_xx_wait_for_offers_xx_xx_xxx_offers.mbx
+    _output_file_pattern = re.compile(
+        r"^.*git.*/(?P<filename>[\w-]+\.mbx).*$", re.MULTILINE
+    )
+
+    @property
+    def command(self) -> str:
+        return "b4"
+
+    @property
+    def can_install(self) -> bool:
+        return True
+
+    def _install(self) -> bool:
+        if isinstance(self.node.os, Debian):
+            self.node.os.install_packages("b4")
+        installed = self._check_exists()
+        if not installed:
+            pip = self.node.tools[Pip]
+            pip.install_packages("b4", install_to_user=True)
+        return self._check_exists()
+
+    def am(
+        self, message_id: str, output_dir: pathlib.PurePath, sudo: bool = False
+    ) -> pathlib.PurePath:
+        result = self.run(
+            f"am -o '{output_dir}' '{message_id}'",
+            force_run=True,
+            expected_exit_code=0,
+            sudo=sudo,
+        )
+        filename = find_group_in_lines(
+            lines=result.stdout, pattern=self._output_file_pattern, single_line=False
+        ).get("filename")
+        if not filename:
+            raise LisaException("Failed to get filename from b4 am output")
+        return pathlib.PurePath(output_dir, filename)