From caf4d779e30ad7b051ad1076021b15cbc2711d8d Mon Sep 17 00:00:00 2001
From: hexaheximal <hexaheximal@proton.me>
Date: Wed, 12 Jul 2023 23:18:45 +0000
Subject: [PATCH] Add mtkclient as a flasher option

Since we're using mtkclient a lot anyways, it makes sense to add
support for it to pmbootstrap.

This was originally implemented by JustSoup321, but they had issues
submitting the patch, so I've cleaned it up for upstream submission.

Co-Authored-By: JustSoup321 <brandonboese@protonmail.com>
Co-Authored-By: Oliver Smith <ollieparanoid@postmarketos.org>
Signed-off-by: hexaheximal <hexaheximal@proton.me>
Reviewed-by: Oliver Smith <ollieparanoid@postmarketos.org>
Link: https://lists.sr.ht/~postmarketos/pmbootstrap-devel/%3C20230712231807.22590-1-hexaheximal@proton.me%3E
---
 pmb/config/__init__.py   | 26 ++++++++++++++++++++++++++
 pmb/flasher/init.py      |  4 ++++
 pmb/flasher/variables.py |  9 +++++++++
 3 files changed, 39 insertions(+)

diff --git a/pmb/config/__init__.py b/pmb/config/__init__.py
index 6feb73dd..df8bbdcc 100644
--- a/pmb/config/__init__.py
+++ b/pmb/config/__init__.py
@@ -802,6 +802,10 @@ deviceinfo_attributes = [
     "flash_rk_partition_kernel",
     "flash_rk_partition_rootfs",
     "flash_rk_partition_system", # deprecated
+    "flash_mtkclient_partition_kernel",
+    "flash_mtkclient_partition_rootfs",
+    "flash_mtkclient_partition_vbmeta",
+    "flash_mtkclient_partition_dtbo",
     "generate_legacy_uboot_initfs",
     "kernel_cmdline",
     "generate_bootimg",
@@ -881,6 +885,7 @@ flash_methods = [
     "0xffff",
     "fastboot",
     "heimdall",
+    "mtkclient",
     "none",
     "rkdeveloptool",
     "uuu",
@@ -1031,6 +1036,27 @@ flashers = {
                     "$IMAGE_SPLIT_BOOT"]
             ],
         },
+    },
+    "mtkclient": {
+        "depends": ["mtkclient"],
+        "actions": {
+            "flash_rootfs": [["mtk", "w", "$PARTITION_ROOTFS",
+                              "$IMAGE"]],
+            "flash_kernel": [["mtk", "w", "$PARTITION_KERNEL",
+                              "$BOOT/boot.img$FLAVOR"]],
+            "flash_vbmeta": [
+                # Generate vbmeta image with "disable verification" flag
+                ["avbtool", "make_vbmeta_image", "--flags", "2",
+                    "--padding_size", "$FLASH_PAGESIZE",
+                    "--output", "/vbmeta.img"],
+                ["mtk", "w", "$PARTITION_VBMETA", "/vbmeta.img"],
+                ["rm", "-f", "/vbmeta.img"]
+            ],
+            "flash_dtbo": [["mtk", "w", "$PARTITION_DTBO",
+                            "$BOOT/dtbo.img"]],
+            "flash_lk2nd": [["mtk", "w", "$PARTITION_KERNEL",
+                             "$BOOT/lk2nd.img"]]
+        }
     }
 }
 
diff --git a/pmb/flasher/init.py b/pmb/flasher/init.py
index aff7808d..3569bcce 100644
--- a/pmb/flasher/init.py
+++ b/pmb/flasher/init.py
@@ -32,6 +32,10 @@ def install_depends(args):
         pmaports_cfg = pmb.config.pmaports.read_config(args)
         depends = pmaports_cfg.get("supported_heimdall_depends",
                                    "heimdall,avbtool").split(",")
+    elif method == "mtkclient":
+        pmaports_cfg = pmb.config.pmaports.read_config(args)
+        depends = pmaports_cfg.get("supported_mtkclient_depends",
+                                   "mtkclient,android-tools").split(",")
 
     pmb.chroot.apk.install(args, depends)
 
diff --git a/pmb/flasher/variables.py b/pmb/flasher/variables.py
index e00ed1d7..6fd5469a 100644
--- a/pmb/flasher/variables.py
+++ b/pmb/flasher/variables.py
@@ -31,6 +31,15 @@ def variables(args, flavor, method):
             or args.deviceinfo["flash_rk_partition_system"] or None
         _partition_vbmeta = None
         _partition_dtbo = None
+    elif method.startswith("mtkclient"):
+        _partition_kernel = args.deviceinfo["flash_mtkclient_partition_kernel"]\
+            or "boot"
+        _partition_rootfs = args.deviceinfo["flash_mtkclient_partition_rootfs"]\
+            or "userdata"
+        _partition_vbmeta = args.deviceinfo["flash_mtkclient_partition_vbmeta"]\
+            or None
+        _partition_dtbo = args.deviceinfo["flash_mtkclient_partition_dtbo"]\
+            or None
     else:
         _partition_kernel = args.deviceinfo["flash_heimdall_partition_kernel"]\
             or "KERNEL"