From f338168caba2b100eca85e5e8dfcea78cb27f1e2 Mon Sep 17 00:00:00 2001
From: Bjoern Kerler <info@revskills.de>
Date: Sat, 3 Aug 2024 17:28:38 +0200
Subject: [PATCH] Fix reinit and rpmb issues

---
 mtkclient/Library/DA/mtk_da_handler.py          |  3 ++-
 mtkclient/Library/DA/mtk_daloader.py            |  3 +--
 mtkclient/Library/DA/xflash/extension/xflash.py | 15 +++++++++++----
 mtkclient/Library/DA/xflash/xflash_lib.py       |  3 ---
 mtkclient/Library/DA/xml/extension/v6.py        | 11 +++++++----
 mtkclient/Library/mtk_preloader.py              |  5 +++++
 mtkclient/Library/settings.py                   |  4 +++-
 mtkclient/config/mtk_config.py                  |  1 -
 8 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/mtkclient/Library/DA/mtk_da_handler.py b/mtkclient/Library/DA/mtk_da_handler.py
index dfcce12..98b0cec 100755
--- a/mtkclient/Library/DA/mtk_da_handler.py
+++ b/mtkclient/Library/DA/mtk_da_handler.py
@@ -10,6 +10,7 @@ from mtkclient.Library.error import ErrorHandler
 from mtkclient.Library.utils import Progress
 from mtkclient.config.brom_config import Efuse, DAmodes
 from mtkclient.Library.Filesystem.mtkdafs import MtkDaFS
+
 try:
     from fuse import FUSE
 except ImportError:
@@ -30,7 +31,7 @@ class DaHandler(metaclass=LogBase):
         self.pid = mtk.config.pid
         self.interface = mtk.config.interface
         self.pathconfig = PathConfig()
-        self.__logger, self.info, self.debug, self.warning, self.error = logsetup(self, self.__logger, 
+        self.__logger, self.info, self.debug, self.warning, self.error = logsetup(self, self.__logger,
                                                                                   loglevel, mtk.config.gui)
         self.eh = ErrorHandler()
         self.mtk = mtk
diff --git a/mtkclient/Library/DA/mtk_daloader.py b/mtkclient/Library/DA/mtk_daloader.py
index d975c00..0e4788d 100755
--- a/mtkclient/Library/DA/mtk_daloader.py
+++ b/mtkclient/Library/DA/mtk_daloader.py
@@ -145,6 +145,7 @@ class DAloader(metaclass=LogBase):
                 self.config.meid = bytes.fromhex(config["meid"])
             if "socid" in config:
                 self.config.socid = bytes.fromhex(config["socid"])
+            self.config.hwparam = HwParam(self.mtk.config, self.config.meid.hex(), self.mtk.config.hwparam_path)
             if config["flashmode"] == "LEGACY":
                 self.mtk.config.chipconfig.damode = DAmodes.LEGACY
                 self.flashmode = DAmodes.LEGACY
@@ -155,8 +156,6 @@ class DAloader(metaclass=LogBase):
                 self.mtk.config.chipconfig.damode = DAmodes.XML
                 self.flashmode = DAmodes.XML
             self.config.init_hwcode(self.config.hwcode)
-            if self.config.meid is not None:
-                self.config.hwparam = HwParam(self.mtk.config, self.config.meid.hex(), self.mtk.config.hwparam_path)
             if self.flashmode == DAmodes.XML:
                 self.da = DAXML(self.mtk, self.daconfig, self.loglevel)
                 self.daconfig.flashtype = config["flashtype"]
diff --git a/mtkclient/Library/DA/xflash/extension/xflash.py b/mtkclient/Library/DA/xflash/extension/xflash.py
index 257ba15..f03eab4 100755
--- a/mtkclient/Library/DA/xflash/extension/xflash.py
+++ b/mtkclient/Library/DA/xflash/extension/xflash.py
@@ -127,6 +127,7 @@ class XFlashExt(metaclass=LogBase):
             ufshcd_queuecommand_ptr = daextdata.find(b"\x55\x55\x55\x55")
             ufshcd_get_free_tag_ptr = daextdata.find(b"\x66\x66\x66\x66")
             ptr_g_ufs_hba_ptr = daextdata.find(b"\x77\x77\x77\x77")
+            efuse_addr_ptr = daextdata.find(b"\x88\x88\x88\x88")
 
             if register_ptr != -1 and mmc_get_card_ptr != -1:
                 if register_devctrl:
@@ -159,6 +160,8 @@ class XFlashExt(metaclass=LogBase):
                 if g_ufs_hba is None:
                     g_ufs_hba = 0
 
+                efuse_addr = self.config.chipconfig.efuse_addr
+
                 # Patch the addr
                 daextdata[register_ptr:register_ptr + 4] = pack("<I", register_devctrl)
                 daextdata[mmc_get_card_ptr:mmc_get_card_ptr + 4] = pack("<I", mmc_get_card)
@@ -167,6 +170,7 @@ class XFlashExt(metaclass=LogBase):
                 daextdata[ufshcd_get_free_tag_ptr:ufshcd_get_free_tag_ptr + 4] = pack("<I", ufshcd_get_free_tag)
                 daextdata[ufshcd_queuecommand_ptr:ufshcd_queuecommand_ptr + 4] = pack("<I", ufshcd_queuecommand)
                 daextdata[ptr_g_ufs_hba_ptr:ptr_g_ufs_hba_ptr + 4] = pack("<I", g_ufs_hba)
+                daextdata[efuse_addr_ptr:efuse_addr_ptr+4]=pack("<I", efuse_addr)
 
                 # print(hexlify(daextdata).decode('utf-8'))
                 # open("daext.bin","wb").write(daextdata)
@@ -315,10 +319,10 @@ class XFlashExt(metaclass=LogBase):
     def custom_read(self, addr, length):
         data = bytearray()
         pos = 0
-        while pos<length:
+        while pos < length:
             if self.cmd(XCmd.CUSTOM_READ):
-                self.xsend(data=addr+pos, is64bit=True)
-                sz = min(length,0x10000)
+                self.xsend(data=addr + pos, is64bit=True)
+                sz = min(length, 0x10000)
                 self.xsend(sz)
                 tmp = self.xread()
                 data.extend(tmp)
@@ -443,7 +447,8 @@ class XFlashExt(metaclass=LogBase):
             if meid != b"\x00" * 16:
                 # self.config.set_meid(meid)
                 self.info("Generating sej rpmbkey...")
-                rpmbkey = hwc.aes_hwcrypt(mode="rpmb", data=meid, btype="sej", otp=otp)
+                #rpmbkey = hwc.aes_hwcrypt(mode="rpmb", data=meid, btype="sej", otp=otp)
+                rpmbkey = hwc.aes_hwcrypt(btype="dxcc", mode="rpmb")
                 if rpmbkey is not None:
                     if self.cmd(XCmd.CUSTOM_SET_RPMB_KEY):
                         self.xsend(rpmbkey)
@@ -461,6 +466,8 @@ class XFlashExt(metaclass=LogBase):
             cmd = XCmd.CUSTOM_INIT_UFS_RPMB
         if self.cmd(cmd):
             derivedrpmb = self.xread()
+            # rpmb_frame = self.xread()
+            # hash = self.xread()
             if int.from_bytes(derivedrpmb[:4], 'little') != 0xff:
                 status = self.status()
                 if status == 0:
diff --git a/mtkclient/Library/DA/xflash/xflash_lib.py b/mtkclient/Library/DA/xflash/xflash_lib.py
index f69f10a..f69af04 100755
--- a/mtkclient/Library/DA/xflash/xflash_lib.py
+++ b/mtkclient/Library/DA/xflash/xflash_lib.py
@@ -1245,9 +1245,6 @@ class DAXFlash(metaclass=LogBase):
                             if status == 0x0 and unpack("<I", ret)[0] == 0xA1A2A3A4:
                                 self.info("DA Extensions successfully added")
                                 self.daext = True
-                                self.config.hwparam = HwParam(self.mtk.config, self.mtk.config.meid,
-                                                              self.mtk.config.hwparam_path)
-                                self.config.hwparam.writesetting("hwcode", hex(self.config.hwcode))
                         if not self.daext:
                             self.warning("DA Extensions failed to enable")
 
diff --git a/mtkclient/Library/DA/xml/extension/v6.py b/mtkclient/Library/DA/xml/extension/v6.py
index f73d147..b7a28d3 100755
--- a/mtkclient/Library/DA/xml/extension/v6.py
+++ b/mtkclient/Library/DA/xml/extension/v6.py
@@ -1,3 +1,4 @@
+import hashlib
 import logging
 import os
 import sys
@@ -135,6 +136,7 @@ class XmlFlashExt(metaclass=LogBase):
             ufshcd_queuecommand_ptr = daextdata.find(b"\x55\x55\x55\x55")
             ufshcd_get_free_tag_ptr = daextdata.find(b"\x66\x66\x66\x66")
             ptr_g_ufs_hba_ptr = daextdata.find(b"\x77\x77\x77\x77")
+            efuse_addr_ptr = daextdata.find(b"\x88\x88\x88\x88")
 
             # register_xml_cmd("CMD:GET-SYS-PROPERTY", & a1, cmd_get_sys_property);
 
@@ -181,6 +183,7 @@ class XmlFlashExt(metaclass=LogBase):
             if mmc_rpmb_send_command is None:
                 mmc_rpmb_send_command = 0
 
+            efuse_addr = self.config.chipconfig.efuse_addr
             #########################################
             if register_ptr != -1:
                 if register_xml_cmd:
@@ -196,6 +199,7 @@ class XmlFlashExt(metaclass=LogBase):
                 daextdata[ufshcd_get_free_tag_ptr:ufshcd_get_free_tag_ptr + 4] = pack("<I", ufshcd_get_free_tag)
                 daextdata[ufshcd_queuecommand_ptr:ufshcd_queuecommand_ptr + 4] = pack("<I", ufshcd_queuecommand)
                 daextdata[ptr_g_ufs_hba_ptr:ptr_g_ufs_hba_ptr + 4] = pack("<I", g_ufs_hba)
+                daextdata[efuse_addr_ptr:efuse_addr_ptr + 4] = pack("<I", efuse_addr)
 
                 # print(hexlify(daextdata).decode('utf-8'))
                 # open("daext.bin","wb").write(daextdata)
@@ -358,7 +362,8 @@ class XmlFlashExt(metaclass=LogBase):
             if meid != b"\x00" * 16:
                 # self.config.set_meid(meid)
                 self.info("Generating sej rpmbkey...")
-                rpmbkey = hwc.aes_hwcrypt(mode="rpmb", data=meid, btype="sej", otp=otp)
+                rpmbkey = hwc.aes_hwcrypt(btype="dxcc", mode="rpmb")
+                #rpmbkey = hwc.aes_hwcrypt(mode="rpmb", data=meid, btype="sej", otp=otp)
                 if rpmbkey is not None:
                     xmlcmd = self.xflash.Cmd.create_cmd("CUSTOMRPMBKEY")
                     if self.xsend(xmlcmd):
@@ -375,9 +380,7 @@ class XmlFlashExt(metaclass=LogBase):
                             if rpmbkey == read_key:
                                 self.info("Setting rpmbkey: ok")
         ufs = False
-        if self.xflash.emmc.rpmb_size != 0:
-            ufs = False
-        elif self.xflash.ufs.block_size != 0:
+        if self.xflash.emmc is None:
             ufs = True
         xmlcmd = self.xflash.Cmd.create_cmd("CUSTOMMMCINIT")
         if ufs:
diff --git a/mtkclient/Library/mtk_preloader.py b/mtkclient/Library/mtk_preloader.py
index 2cad03b..d897304 100755
--- a/mtkclient/Library/mtk_preloader.py
+++ b/mtkclient/Library/mtk_preloader.py
@@ -11,6 +11,7 @@ from binascii import hexlify
 from Cryptodome.Util.number import size
 from mtkclient.Library.Auth.sla_keys import brom_sla_keys
 from mtkclient.Library.Auth.sla import generate_brom_sla_challenge
+from mtkclient.Library.settings import HwParam
 from mtkclient.Library.utils import LogBase, logsetup
 from mtkclient.Library.error import ErrorHandler
 from mtkclient.config.brom_config import DAmodes
@@ -224,6 +225,8 @@ class Preloader(metaclass=LogBase):
             self.info("\tSW Ver:\t\t\t" + hex(self.config.swver))
         meid = self.get_meid()
         if meid is not None:
+            self.config.hwparam = HwParam(self.mtk.config, self.config.meid.hex(), self.mtk.config.hwparam_path)
+            self.config.hwparam.writesetting("hwcode", hex(self.config.hwcode))
             self.config.set_meid(meid)
             if self.display:
                 self.info("ME_ID:\t\t\t" + hexlify(meid).decode('utf-8').upper())
@@ -234,6 +237,8 @@ class Preloader(metaclass=LogBase):
                 if self.display:
                     if socid != b"":
                         self.info("SOC_ID:\t\t\t" + hexlify(socid).decode('utf-8').upper())
+                        self.config.hwparam.writesetting("socid", hexlify(socid).decode('utf-8'))
+
         if self.config.auth is not None and self.config.is_brom and self.config.target_config["daa"]:
             if os.path.exists(self.config.auth):
                 authdata = open(self.config.auth, "rb").read()
diff --git a/mtkclient/Library/settings.py b/mtkclient/Library/settings.py
index f09ea82..e8f233c 100755
--- a/mtkclient/Library/settings.py
+++ b/mtkclient/Library/settings.py
@@ -26,7 +26,9 @@ class HwParam:
             self.paramsetting = {}
             if os.path.exists(os.path.join(path, self.paramfile)):
                 try:
-                    self.paramsetting = json.loads(open(os.path.join(path, self.paramfile), "r").read())
+                    tmp = json.loads(open(os.path.join(path, self.paramfile), "r").read())
+                    if tmp["meid"] == meid:
+                        self.paramsetting = tmp
                 except Exception:
                     # json file invalid, load nothing.
                     pass
diff --git a/mtkclient/config/mtk_config.py b/mtkclient/config/mtk_config.py
index 69f83a3..1b37196 100755
--- a/mtkclient/config/mtk_config.py
+++ b/mtkclient/config/mtk_config.py
@@ -112,7 +112,6 @@ class MtkConfig(metaclass=LogBase):
         self.hwparam.writesetting("hwcode", hex(hwcode))
 
     def set_meid(self, meid):
-        self.hwparam = HwParam(self, meid, self.hwparam_path)
         self.meid = meid
         self.hwparam.writesetting("meid", hexlify(meid).decode('utf-8'))