From 48523af093b300b4e438cbafd3542c84c674188b Mon Sep 17 00:00:00 2001 From: bongbui321 Date: Tue, 12 Mar 2024 01:01:58 -0400 Subject: [PATCH] cleanup + faster --- edlclient/Library/firehose.py | 33 ++++++++++++---------- edlclient/Library/gpt.py | 52 ----------------------------------- 2 files changed, 18 insertions(+), 67 deletions(-) diff --git a/edlclient/Library/firehose.py b/edlclient/Library/firehose.py index 3462ed9..64d2f9f 100755 --- a/edlclient/Library/firehose.py +++ b/edlclient/Library/firehose.py @@ -1375,10 +1375,10 @@ class firehose(metaclass=LogBase): slot_b_status = not slot_a_status fpartitions = {} try: - for lun in self.luns: - lunname = "Lun" + str(lun) + for lun_a in self.luns: + lunname = "Lun" + str(lun_a) fpartitions[lunname] = [] - header_data_a, guid_gpt_a = self.get_gpt(lun, int(0), int(0), int(0)) + header_data_a, guid_gpt_a = self.get_gpt(lun_a, int(0), int(0), int(0)) if guid_gpt_a is None: break else: @@ -1386,11 +1386,16 @@ class firehose(metaclass=LogBase): slot = partitionname_a.lower()[-2:] if slot == "_a": partitionname_b = partitionname_a[:-1] + "b" - resp = self.detect_partition(arguments=None, partitionname=partitionname_b, send_full=True) - if not resp[0]: - self.error(f"Cannot find partition {partitionname_b}") - return False - _, lun_b, header_data_b, guid_gpt_b = resp + if partitionname_b in guid_gpt_a.partentries: + lun_b = lun_a + header_data_b = header_data_a + guid_gpt_b = guid_gpt_a + else: + resp = self.detect_partition(arguments=None, partitionname=partitionname_b, send_full=True) + if not resp[0]: + self.error(f"Cannot find partition {partitionname_b}") + return False + _, lun_b, header_data_b, guid_gpt_b = resp part_a = guid_gpt_a.partentries[partitionname_a] part_b = guid_gpt_b.partentries[partitionname_b] @@ -1403,26 +1408,24 @@ class firehose(metaclass=LogBase): slot_a_status, slot_b_status, is_boot ) + assert(poffset_b != poffset_a) + header_data_a[poffset_a : poffset_a+len(pdata_a)] = pdata_a new_header_a = guid_gpt_a.fix_gpt_crc(header_data_a) - - if lun == lun_b: - assert(poffset_b != poffset_a) - header_data_b = new_header_a header_data_b[poffset_b:poffset_b + len(pdata_b)] = pdata_b new_header_b = guid_gpt_b.fix_gpt_crc(header_data_b) if new_header_a is not None: start_sector_patch_a = poffset_a // self.cfg.SECTOR_SIZE_IN_BYTES byte_offset_patch_a = poffset_a % self.cfg.SECTOR_SIZE_IN_BYTES - cmd_patch_multiple(lun, start_sector_patch_a, byte_offset_patch_a, pdata_a) + cmd_patch_multiple(lun_a, start_sector_patch_a, byte_offset_patch_a, pdata_a) # header will be updated in partitionname_b if in same lun - if lun != lun_b: + if lun_a != lun_b: headeroffset_a = guid_gpt_a.header.current_lba * guid_gpt_a.sectorsize start_sector_hdr_a = guid_gpt_a.header.current_lba pheader_a = new_header_a[headeroffset_a : headeroffset_a+guid_gpt_a.header.header_size] - cmd_patch_multiple(lun, start_sector_hdr_a, 0, pheader_a) + cmd_patch_multiple(lun_a, start_sector_hdr_a, 0, pheader_a) if new_header_b is not None: start_sector_patch_b = poffset_b // self.cfg.SECTOR_SIZE_IN_BYTES diff --git a/edlclient/Library/gpt.py b/edlclient/Library/gpt.py index b3b8416..070819c 100755 --- a/edlclient/Library/gpt.py +++ b/edlclient/Library/gpt.py @@ -496,58 +496,6 @@ class gpt(metaclass=LogBase): res = self.print_gptfile(os.path.join("TestFiles", "gpt_sm8180x.bin")) assert res, "GPT Partition wasn't decoded properly" - def patch(self, data:bytes, partitionname="boot", active: bool = True): - def set_flags(flags, active): - new_flags = flags - if active: - new_flags |= AB_PARTITION_ATTR_SLOT_ACTIVE << (AB_FLAG_OFFSET*8) - else: - new_flags &= ~(AB_PARTITION_ATTR_SLOT_ACTIVE << (AB_FLAG_OFFSET*8)) - return new_flags - try: - rf = BytesIO(data) - for sectorsize in [512, 4096]: - result = self.parse(data, sectorsize) - if result: - for rname_a in self.partentries: - if partitionname.lower() == rname_a.lower(): - rname_b = rname_a[:-1] + "b" - if rname_b not in self.partentries: - return None, None, None, None - - slot_a_state = active - slot_b_state = not active - - partition_a = self.partentries[rname_a] - partition_b = self.partentries[rname_b] - - rf.seek(partition_a.entryoffset) - sdata_a = rf.read(self.header.part_entry_size) - partentry_a = self.gpt_partition(sdata_a) - partentry_a.flags = set_flags(partentry_a.flags, slot_a_state) - - rf.seek(partition_b.entryoffset) - sdata_b = rf.read(self.header.part_entry_size) - partentry_b = self.gpt_partition(sdata_b) - partentry_b.flags = set_flags(partentry_b.flags, slot_b_state) - - print(f"[Before] pa: {partentry_a.type}, pb: {partentry_b.type}") - old_a = partentry_a.type - old_b = partentry_b.type - partentry_a.type, partentry_b.type = partentry_b.type, partentry_a.type - assert(partentry_a.type == old_b) - assert(partentry_b.type == old_a) - print(f"[After] pa: {partentry_a.type}, pb: {partentry_b.type}") - - pdata_a = partentry_a.create() - pdata_b = partentry_b.create() - return pdata_a, pdata_b, partition_a.entryoffset, partition_b.entryoffset - break - return None, None, None, None - except Exception as e: - self.error(str(e)) - return None, None, None, None - def fix_gpt_crc(self, data): partentry_size = self.header.num_part_entries * self.header.part_entry_size partentry_offset = self.header.part_entry_start_lba * self.sectorsize