From c51877d588d9be3da0b8fc5e07cc4605c8367005 Mon Sep 17 00:00:00 2001
From: CTCaer <ctcaer@gmail.com>
Date: Sun, 8 May 2022 05:59:32 +0300
Subject: [PATCH] nyx: tools: add mbr hidden attr fixer in hybrid mbr fixer

---
 .../frontend/gui_tools_partition_manager.c    | 51 ++++++++++++++++---
 1 file changed, 45 insertions(+), 6 deletions(-)

diff --git a/nyx/nyx_gui/frontend/gui_tools_partition_manager.c b/nyx/nyx_gui/frontend/gui_tools_partition_manager.c
index 344cc42..5763921 100644
--- a/nyx/nyx_gui/frontend/gui_tools_partition_manager.c
+++ b/nyx/nyx_gui/frontend/gui_tools_partition_manager.c
@@ -2058,6 +2058,7 @@ static lv_res_t _action_fix_mbr(lv_obj_t *btn)
 
 	mbr_t mbr[2] = { 0 };
 	gpt_t *gpt = calloc(1, sizeof(gpt_t));
+	gpt_header_t gpt_hdr_backup = { 0 };
 
 	sdmmc_storage_read(&sd_storage, 0, 1, &mbr[0]);
 	sdmmc_storage_read(&sd_storage, 1, sizeof(gpt_t) >> 9, gpt);
@@ -2069,9 +2070,12 @@ static lv_res_t _action_fix_mbr(lv_obj_t *btn)
 	if (memcmp(&gpt->header.signature, "EFI PART", 8) || gpt->header.num_part_ents > 128)
 	{
 		lv_label_set_text(lbl_status, "#FFDD00 Warning:# No valid GPT was found!");
+		free(gpt);
 		goto out;
 	}
 
+	sdmmc_storage_read(&sd_storage, gpt->header.alt_lba, 1, &gpt_hdr_backup);
+
 	// Parse GPT.
 	LIST_INIT(gpt_parsed);
 	for (u32 i = 0; i < gpt->header.num_part_ents; i++)
@@ -2092,7 +2096,6 @@ static lv_res_t _action_fix_mbr(lv_obj_t *btn)
 
 		list_append(&gpt_parsed, &part->link);
 	}
-	free(gpt);
 
 	// Set FAT and emuMMC partitions.
 	u32 mbr_idx = 1;
@@ -2138,21 +2141,27 @@ static lv_res_t _action_fix_mbr(lv_obj_t *btn)
 	mbr[1].partitions[mbr_idx].size_sct = sd_storage.sec_cnt - 1;
 
 	// Check for differences.
-	bool changed = false;
+	bool hybrid_mbr_changed = false;
 	for (u32 i = 1; i < 4; i++)
 	{
 		if ((mbr[0].partitions[i].type      != mbr[1].partitions[i].type)      ||
 			(mbr[0].partitions[i].start_sct != mbr[1].partitions[i].start_sct) ||
 			(mbr[0].partitions[i].size_sct  != mbr[1].partitions[i].size_sct))
 		{
-			changed = true;
+			hybrid_mbr_changed = true;
 			break;
 		}
 	}
 
-	if (!changed)
+	// Check for secret MBR attributes.
+	bool has_mbr_attributes = false;
+	if (gpt->entries[0].part_guid[7])
+		has_mbr_attributes = true;
+
+	if (!hybrid_mbr_changed && !has_mbr_attributes)
 	{
 		lv_label_set_text(lbl_status, "#96FF00 Warning:# The Hybrid MBR needs no change!#");
+		free(gpt);
 		goto out;
 	}
 
@@ -2200,14 +2209,44 @@ static lv_res_t _action_fix_mbr(lv_obj_t *btn)
 
 	if (btn_wait() & BTN_POWER)
 	{
-		// Write MBR.
 		sd_mount();
-		sdmmc_storage_write(&sd_storage, 0, 1, &mbr[1]);
+
+		// Write MBR.
+		if (hybrid_mbr_changed)
+			sdmmc_storage_write(&sd_storage, 0, 1, &mbr[1]);
+
+		// Fix MBR secret attributes.
+		if (has_mbr_attributes)
+		{
+			gpt->entries[0].part_guid[7] = 0;
+
+			u32 entries_size = sizeof(gpt_entry_t) * gpt->header.num_part_ents;
+			gpt->header.part_ents_crc32 = crc32_calc(0, (const u8 *)gpt->entries, entries_size);
+			gpt->header.crc32 = 0; // Set to 0 for calculation.
+			gpt->header.crc32 = crc32_calc(0, (const u8 *)&gpt->header, gpt->header.size);
+
+			gpt_hdr_backup.part_ents_crc32 = gpt->header.part_ents_crc32;
+			gpt_hdr_backup.crc32 = 0; // Set to 0 for calculation.
+			gpt_hdr_backup.crc32 = crc32_calc(0, (const u8 *)&gpt_hdr_backup, gpt_hdr_backup.size);
+
+			// Write main GPT.
+			u32 aligned_entries_size = ALIGN(entries_size, 512);
+			sdmmc_storage_write(&sd_storage, gpt->header.my_lba, (sizeof(gpt_header_t) + aligned_entries_size) >> 9, gpt);
+
+			// Write backup GPT partition table.
+			sdmmc_storage_write(&sd_storage, gpt_hdr_backup.part_ent_lba, aligned_entries_size >> 9, gpt->entries);
+
+			// Write backup GPT header.
+			sdmmc_storage_write(&sd_storage, gpt_hdr_backup.my_lba, 1, &gpt_hdr_backup);
+		}
+
 		sd_unmount();
+
 		lv_label_set_text(lbl_status, "#96FF00 The new Hybrid MBR was written successfully!#");
 	}
 	else
 		lv_label_set_text(lbl_status, "#FFDD00 Warning: The Hybrid MBR Fix was canceled!#");
+	free(gpt);
 
 out:
 	lv_mbox_add_btns(mbox, mbox_btn_map, mbox_action);