From 8c6183e061962ed719c68cae08036920e73e581e Mon Sep 17 00:00:00 2001
From: Tobias Sterbak <sterbak-it@outlook.com>
Date: Fri, 13 Oct 2023 12:28:21 +0200
Subject: [PATCH] Change fastboot_flash_recovery to also flash additional
 partitions if nessecary (only for Mi439)

---
 .../assets/configs/Mi439.yaml                 |  8 +-----
 openandroidinstaller/assets/configs/vayu.yaml |  2 +-
 openandroidinstaller/tooling.py               | 28 +++++++++++++++++--
 openandroidinstaller/views/step_view.py       |  4 +++
 4 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/openandroidinstaller/assets/configs/Mi439.yaml b/openandroidinstaller/assets/configs/Mi439.yaml
index cc286a80..20b4afe0 100644
--- a/openandroidinstaller/assets/configs/Mi439.yaml
+++ b/openandroidinstaller/assets/configs/Mi439.yaml
@@ -49,18 +49,12 @@ steps:
         
         Once the device is fully booted, you need to reboot into the bootloader again by pressing 'Confirm and run' here. Then continue.
       command: adb_reboot_bootloader
-    - type: call_button
-      content: >
-        Install additional partitions selected before by pressing 'Confirm and run'. Once it's done continue.
-        
-        Note : If you have not selected additional partitions it will do nothing and is save to continue.
-      command: fastboot_flash_additional_partitions
     - type: call_button
       content: >
         Install the recovery you chosen before by pressing 'Confirm and run'. Once it's done continue.
       command: fastboot_flash_recovery
     - type: call_button
-      img: ofox.png
+      img: twrp-start.jpeg
       content: >
         Reboot to recovery by pressing 'Confirm and run', and hold the Vol+ button of your phone UNTIL you see the recovery.
         If MiUI starts, you have to start the process again, since MiUI delete the recovery you just flashed.
diff --git a/openandroidinstaller/assets/configs/vayu.yaml b/openandroidinstaller/assets/configs/vayu.yaml
index 16223090..f040e367 100644
--- a/openandroidinstaller/assets/configs/vayu.yaml
+++ b/openandroidinstaller/assets/configs/vayu.yaml
@@ -46,7 +46,7 @@ steps:
         Install the recovery you chosen before by pressing 'Confirm and run'. Once it's done continue.
       command: fastboot_flash_recovery
     - type: call_button
-      img: ofox.png
+      img: twrp-start.jpeg
       content: >
         Reboot to recovery by pressing 'Confirm and run', and hold the Vol+ button of your phone UNTIL you see the recovery.
         If MiUI starts, you have to start the process again, since MiUI delete the recovery you just flashed.
diff --git a/openandroidinstaller/tooling.py b/openandroidinstaller/tooling.py
index b82a3276..69cfd0a7 100644
--- a/openandroidinstaller/tooling.py
+++ b/openandroidinstaller/tooling.py
@@ -451,9 +451,29 @@ def fastboot_flash_boot(bin_path: Path, recovery: str) -> TerminalResponse:
 
 @add_logging("Flash custom recovery with fastboot.")
 def fastboot_flash_recovery(
-    bin_path: Path, recovery: str, is_ab: bool = True
+    bin_path: Path,
+    recovery: str,
+    is_ab: bool = True,
+    vendor_boot: Optional[str] = None,
+    dtbo: Optional[str] = None,
+    vbmeta: Optional[str] = None,
+    super_empty: Optional[str] = None,
 ) -> TerminalResponse:
-    """Flash custom recovery with fastboot."""
+    """Flash custom recovery with fastboot.
+
+    If necessary, flash additional partitions (dtbo, vbmeta, super_empty) with fastboot before.
+    """
+    if any([dtbo, vbmeta, super_empty, vendor_boot]):
+        for line in fastboot_flash_additional_partitions(
+            bin_path=bin_path,
+            dtbo=dtbo,
+            vbmeta=vbmeta,
+            super_empty=super_empty,
+            vendor_boot=vendor_boot,
+            is_ab=is_ab,
+        ):
+            yield line
+
     for line in run_command(
         "fastboot flash recovery ", target=f"{recovery}", bin_path=bin_path
     ):
@@ -505,7 +525,9 @@ def fastboot_flash_additional_partitions(
     if vbmeta:
         logger.info("vbmeta selected. Flashing vbmeta partition.")
         for line in run_command(
-            "fastboot --disable-verity --disable-verification flash vbmeta ", target=f"{vbmeta}", bin_path=bin_path
+            "fastboot --disable-verity --disable-verification flash vbmeta ",
+            target=f"{vbmeta}",
+            bin_path=bin_path,
         ):
             yield line
         if not is_ab:
diff --git a/openandroidinstaller/views/step_view.py b/openandroidinstaller/views/step_view.py
index ef467351..b64a78c2 100644
--- a/openandroidinstaller/views/step_view.py
+++ b/openandroidinstaller/views/step_view.py
@@ -242,6 +242,10 @@ def call_to_phone(self, e, command: str):
                 fastboot_flash_recovery,
                 recovery=self.state.recovery_path,
                 is_ab=self.state.config.is_ab,
+                dtbo=self.state.dtbo_path,
+                vbmeta=self.state.vbmeta_path,
+                super_empty=self.state.super_empty_path,
+                vendor_boot=self.state.vendor_boot_path,
             ),
             "fastboot_reboot_recovery": fastboot_reboot_recovery,
             "fastboot_flash_additional_partitions": partial(