From 4ebe85afb65704d6c4651b312c6d3a9659cb139c Mon Sep 17 00:00:00 2001
From: BigfootACA <bigfoot@classfun.cn>
Date: Tue, 21 May 2024 11:21:48 +0800
Subject: [PATCH] builder: disk: layout: gpt: add initial hybrid support

Signed-off-by: BigfootACA <bigfoot@classfun.cn>
---
 builder/disk/layout/gpt/layout.py | 13 ++++++++++++-
 builder/disk/layout/gpt/part.py   | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/builder/disk/layout/gpt/layout.py b/builder/disk/layout/gpt/layout.py
index 8fc664f..ea113e1 100644
--- a/builder/disk/layout/gpt/layout.py
+++ b/builder/disk/layout/gpt/layout.py
@@ -100,6 +100,10 @@ class DiskLayoutGPT(DiskLayout):
 		puuid = UUID(config["puuid"]) if "puuid" in config else None
 		part = self.add_partition(ptype, area=area, name=pname, uuid=puuid)
 		if part:
+			if "hybrid" in config:
+				part.hybrid = True
+			if "bootable" in config:
+				part.bootable = True
 			if "attributes" in config:
 				part.attributes = config["attributes"]
 		return part
@@ -248,13 +252,20 @@ class DiskLayoutGPT(DiskLayout):
 			size = MasterBootRecord.boot_code.size
 			code = bytes_pad(self.boot_code, size, trunc=True)
 			ctypes.memmove(new_pmbr.boot_code, code, size)
+		idx = 0
+		for part in self.partitions:
+			if not part.hybrid: continue
+			if idx >= 3: raise RuntimeError("Hybrid partition to many")
+			ppart = part.to_mbr_entry()
+			new_pmbr.partitions[idx] = ppart
+			idx += 1
 		ppart = MbrPartEntry()
 		ppart.start_lba = 1
 		ppart.size_lba = self.total_lba - 1
 		ppart.start_head, ppart.start_track, ppart.start_sector = 0, 0, 2
 		ppart.end_head, ppart.end_track, ppart.end_sector = 255, 255, 255
 		ppart.set_type("gpt")
-		new_pmbr.partitions[0] = ppart
+		new_pmbr.partitions[idx] = ppart
 		return new_pmbr
 
 	def create_gpt_entries(self) -> bytes:
diff --git a/builder/disk/layout/gpt/part.py b/builder/disk/layout/gpt/part.py
index f1cc8b9..b6471f9 100644
--- a/builder/disk/layout/gpt/part.py
+++ b/builder/disk/layout/gpt/part.py
@@ -2,15 +2,20 @@ from uuid import UUID
 from logging import getLogger
 from builder.disk.layout.layout import DiskLayout, DiskPart
 from builder.disk.layout.gpt.struct import EfiPartEntry
+from builder.disk.layout.mbr.struct import MbrPartEntry
 from builder.disk.layout.gpt.types import DiskTypesGPT
+from builder.disk.layout.mbr.types import DiskTypesMBR
 from builder.disk.layout.gpt.uefi import EfiGUID
 log = getLogger(__name__)
 
 
 class DiskPartGPT(DiskPart):
 	layout: DiskLayout
+	mbr_type: int
 	type_uuid: UUID
+	bootable: bool
 	uuid: UUID
+	hybrid: bool
 	idx: int
 	_part_name: str
 	_attributes: int
@@ -40,6 +45,10 @@ class DiskPartGPT(DiskPart):
 		tid = DiskTypesGPT.lookup_one_uuid(val)
 		if tid is None: raise ValueError(f"unknown type {val}")
 		self.type_uuid = tid
+		if self.hybrid:
+			tid = DiskTypesMBR.lookup_one_id(val)
+			if tid is None: raise ValueError(f"unknown type {val}")
+			self.mbr_type = tid
 
 	@property
 	def start_lba(self) -> int:
@@ -99,6 +108,14 @@ class DiskPartGPT(DiskPart):
 		part.set_part_name(self.part_name)
 		return part
 
+	def to_mbr_entry(self) -> MbrPartEntry:
+		part = MbrPartEntry()
+		part.set_bootable(self.bootable)
+		part.set_type(self.mbr_type)
+		part.set_start_lba(self.start_lba)
+		part.set_size_lba(self.size_lba)
+		return part
+
 	def __init__(
 		self,
 		layout: DiskLayout,
@@ -108,7 +125,10 @@ class DiskPartGPT(DiskPart):
 		super().__init__()
 		self.layout = layout
 		self.idx = idx
+		self.hybrid = False
+		self.bootable = False
 		self.part_name = None
+		self.mbr_type = 0
 		self.start_lba = 0
 		self.end_lba = 0
 		self.attributes = 0