From a2f362127f7dd333a42af83411381c28d307fb6e Mon Sep 17 00:00:00 2001
From: Tobias Sterbak <sterbak-it@outlook.com>
Date: Sat, 7 Oct 2023 11:27:23 +0200
Subject: [PATCH] Enable flashing of vendor_boot partition

---
 openandroidinstaller/app_state.py         |  1 +
 openandroidinstaller/views/select_view.py | 52 ++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/openandroidinstaller/app_state.py b/openandroidinstaller/app_state.py
index 30f68b08..22d2722f 100644
--- a/openandroidinstaller/app_state.py
+++ b/openandroidinstaller/app_state.py
@@ -52,6 +52,7 @@ def __init__(
         self.dtbo_path = None
         self.vbmeta_path = None
         self.super_empty_path = None
+        self.vendor_boot_path = None
 
         # store views
         self.default_views: List = []
diff --git a/openandroidinstaller/views/select_view.py b/openandroidinstaller/views/select_view.py
index 5c9c116d..45322291 100644
--- a/openandroidinstaller/views/select_view.py
+++ b/openandroidinstaller/views/select_view.py
@@ -106,6 +106,9 @@ def init_visuals(
         self.pick_super_empty_dialog = FilePicker(
             on_result=self.pick_super_empty_result
         )
+        self.pick_vendor_boot_dialog = FilePicker(
+            on_result=self.pick_vendor_boot_result
+        )
 
         self.selected_image = Text("Selected image: ")
         self.selected_recovery = Text("Selected recovery: ")
@@ -118,6 +121,9 @@ def init_visuals(
         self.selected_super_empty = Checkbox(
             fill_color=colors.RED, value=None, disabled=True, tristate=True
         )
+        self.selected_vendor_boot = Checkbox(
+            fill_color=colors.RED, value=None, disabled=True, tristate=True
+        )
 
         # initialize and manage button state.
         self.confirm_button = confirm_button(self.on_confirm)
@@ -131,6 +137,7 @@ def init_visuals(
         self.pick_dtbo_dialog.on_result = self.enable_button_if_ready
         self.pick_vbmeta_dialog.on_result = self.enable_button_if_ready
         self.pick_super_empty_dialog.on_result = self.enable_button_if_ready
+        self.pick_vendor_boot_dialog.on_result = self.enable_button_if_ready
         # back button
         self.back_button = ElevatedButton(
             "Back",
@@ -155,6 +162,7 @@ def build(self):
                 self.pick_dtbo_dialog,
                 self.pick_vbmeta_dialog,
                 self.pick_super_empty_dialog,
+                self.pick_vendor_boot_dialog,
             ]
         )
 
@@ -462,6 +470,27 @@ def toggle_additional_image_selection(self):
                             self.selected_super_empty,
                         ]
                     ),
+                ]
+            )
+        if "vendor_boot" in self.state.config.additional_steps:
+            self.selected_vendor_boot.value = False
+            additional_image_selection.extend(
+                [
+                    Row(
+                        [
+                            FilledButton(
+                                "Pick `vendor_boot.img` image",
+                                icon=icons.UPLOAD_FILE,
+                                on_click=lambda _: self.pick_vendor_boot_dialog.pick_files(
+                                    allow_multiple=False,
+                                    file_type="custom",
+                                    allowed_extensions=["img"],
+                                ),
+                                expand=True,
+                            ),
+                            self.selected_vendor_boot,
+                        ]
+                    ),
                     Divider(),
                 ]
             )
@@ -600,6 +629,26 @@ def pick_super_empty_result(self, e: FilePickerResultEvent):
         # update
         self.selected_super_empty.update()
 
+    def pick_vendor_boot_result(self, e: FilePickerResultEvent):
+        path = ", ".join(map(lambda f: f.name, e.files)) if e.files else "Cancelled!"
+        # update the textfield with the name of the file
+        if e.files:
+            # check if the super_empty works with the device and show the filename in different colors accordingly
+            if path == "vendor_boot.img":
+                self.selected_vendor_boot.fill_color = colors.GREEN
+                self.selected_vendor_boot.value = True
+                self.state.vendor_boot_path = e.files[0].path
+                logger.info(
+                    f"Selected vendor_boot.img from {self.state.vendor_boot_path}"
+                )
+            else:
+                self.selected_vendor_boot.fill_color = colors.RED
+                self.selected_vendor_boot.value = False
+        else:
+            logger.info("No image selected.")
+        # update
+        self.selected_vendor_boot.update()
+
     def enable_button_if_ready(self, e):
         """Enable the confirm button if both files have been selected."""
         if (".zip" in self.selected_image.value) and (
@@ -642,6 +691,8 @@ def enable_button_if_ready(self, e):
                     or "vbmeta" not in self.state.config.additional_steps,
                     self.selected_super_empty.value
                     or "super_empty" not in self.state.config.additional_steps,
+                    self.selected_vendor_boot.value
+                    or "vendor_boot" not in self.state.config.additional_steps,
                 ]
             ):
                 logger.error(
@@ -693,4 +744,3 @@ def enable_button_if_ready(self, e):
             self.right_view.update()
         else:
             self.confirm_button.disabled = True
-            # self.continue_eitherway_button.disabled = True