From 3dcd2ad15f95f021247c9c40e36a5c48400cfd66 Mon Sep 17 00:00:00 2001
From: CTCaer <ctcaer@gmail.com>
Date: Sun, 14 Jun 2020 02:23:02 +0300
Subject: [PATCH] hos: Header for pkg2 is now more proper

---
 bootloader/hos/hos.c   |  4 ++--
 bootloader/hos/pkg2.c  | 21 +++++++++++++++++++--
 bootloader/hos/pkg2.h  |  5 +++--
 nyx/nyx_gui/hos/pkg2.h |  3 ++-
 4 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/bootloader/hos/hos.c b/bootloader/hos/hos.c
index 40e5a49..f712efb 100644
--- a/bootloader/hos/hos.c
+++ b/bootloader/hos/hos.c
@@ -864,7 +864,7 @@ int hos_launch(ini_sec_t *cfg)
 	}
 
 	// Rebuild and encrypt package2.
-	pkg2_build_encrypt((void *)PKG2_LOAD_ADDR, ctxt.kernel, ctxt.kernel_size, &kip1_info, ctxt.new_pkg2);
+	pkg2_build_encrypt((void *)PKG2_LOAD_ADDR, ctxt.kernel, ctxt.kernel_size, &kip1_info, ctxt.new_pkg2, kb);
 
 	gfx_puts("Rebuilt & loaded pkg2\n");
 
@@ -885,7 +885,7 @@ int hos_launch(ini_sec_t *cfg)
 		PMC(APBDEV_PMC_SECURE_SCRATCH32) = 0x104; // Warmboot 3.0.1/.2 PA address id.
 
 	// Finalize per firmware key access.
-	switch (ctxt.pkg1_id->kb)
+	switch (kb)
 	{
 	case KB_FIRMWARE_VERSION_100_200:
 	case KB_FIRMWARE_VERSION_300:
diff --git a/bootloader/hos/pkg2.c b/bootloader/hos/pkg2.c
index 490ce36..4723bd3 100644
--- a/bootloader/hos/pkg2.c
+++ b/bootloader/hos/pkg2.c
@@ -1395,7 +1395,7 @@ DPRINTF("adding kip1 '%s' @ %08X (%08X)\n", ki->kip1->name, (u32)ki->kip1, ki->s
 	return ini1_size;
 }
 
-void pkg2_build_encrypt(void *dst, void *kernel, u32 kernel_size, link_t *kips_info, bool new_pkg2)
+void pkg2_build_encrypt(void *dst, void *kernel, u32 kernel_size, link_t *kips_info, bool new_pkg2, u8 kb)
 {
 	u8 *pdst = (u8 *)dst;
 
@@ -1406,14 +1406,20 @@ void pkg2_build_encrypt(void *dst, void *kernel, u32 kernel_size, link_t *kips_i
 	// Header.
 	pkg2_hdr_t *hdr = (pkg2_hdr_t *)pdst;
 	memset(hdr, 0, sizeof(pkg2_hdr_t));
-	pdst += sizeof(pkg2_hdr_t);
+
+	// Set initial header values.
 	hdr->magic = PKG2_MAGIC;
+	hdr->bl_ver = 0;
+	hdr->pkg2_ver = 0xFF;
+
 	if (!new_pkg2)
 		hdr->base = 0x10000000;
 	else
 		hdr->base = 0x60000;
 DPRINTF("kernel @ %08X (%08X)\n", (u32)kernel, kernel_size);
 
+	pdst += sizeof(pkg2_hdr_t);
+
 	// Kernel.
 	memcpy(pdst, kernel, kernel_size);
 	if (!new_pkg2)
@@ -1436,9 +1442,20 @@ DPRINTF("kernel encrypted\n");
 		ini1_size = _pkg2_ini1_build(pdst, hdr, kips_info, new_pkg2);
 DPRINTF("INI1 encrypted\n");
 
+	// Calculate SHA256 over encrypted Kernel and INI1.
+	u8 *pk2_hash_data = (u8 *)dst + 0x100 + sizeof(pkg2_hdr_t);
+	se_calc_sha256_oneshot(&hdr->sec_sha256[0x20 * PKG2_SEC_KERNEL],
+		(void *)pk2_hash_data, hdr->sec_size[PKG2_SEC_KERNEL]);
+	pk2_hash_data += hdr->sec_size[PKG2_SEC_KERNEL];
+	se_calc_sha256_oneshot(&hdr->sec_sha256[0x20 * PKG2_SEC_INI1],
+		(void *)pk2_hash_data, hdr->sec_size[PKG2_SEC_INI1]);
+
 	//Encrypt header.
+	u8 key_ver = kb ? kb + 1 : 0;
 	*(u32 *)hdr->ctr = 0x100 + sizeof(pkg2_hdr_t) + kernel_size + ini1_size;
+	hdr->ctr[4] = key_ver;
 	se_aes_crypt_ctr(8, hdr, sizeof(pkg2_hdr_t), hdr, sizeof(pkg2_hdr_t), hdr);
 	memset(hdr->ctr, 0 , 0x10);
 	*(u32 *)hdr->ctr = 0x100 + sizeof(pkg2_hdr_t) + kernel_size + ini1_size;
+	hdr->ctr[4] = key_ver;
 }
diff --git a/bootloader/hos/pkg2.h b/bootloader/hos/pkg2.h
index 3360247..ab1ba3c 100644
--- a/bootloader/hos/pkg2.h
+++ b/bootloader/hos/pkg2.h
@@ -68,7 +68,8 @@ typedef struct _pkg2_hdr_t
 	u32 magic;
 	u32 base;
 	u32 pad0;
-	u16 version;
+	u8  pkg2_ver;
+	u8  bl_ver;
 	u16 pad1;
 	u32 sec_size[4];
 	u32 sec_off[4];
@@ -154,6 +155,6 @@ const char* pkg2_patch_kips(link_t *info, char* patchNames);
 
 const pkg2_kernel_id_t *pkg2_identify(u8 *hash);
 pkg2_hdr_t *pkg2_decrypt(void *data, u8 kb);
-void pkg2_build_encrypt(void *dst, void *kernel, u32 kernel_size, link_t *kips_info, bool new_pkg2);
+void pkg2_build_encrypt(void *dst, void *kernel, u32 kernel_size, link_t *kips_info, bool new_pkg2, u8 kb);
 
 #endif
diff --git a/nyx/nyx_gui/hos/pkg2.h b/nyx/nyx_gui/hos/pkg2.h
index ffe97ac..6080e2e 100644
--- a/nyx/nyx_gui/hos/pkg2.h
+++ b/nyx/nyx_gui/hos/pkg2.h
@@ -41,7 +41,8 @@ typedef struct _pkg2_hdr_t
 	u32 magic;
 	u32 base;
 	u32 pad0;
-	u16 version;
+	u8  pkg2_ver;
+	u8  bl_ver;
 	u16 pad1;
 	u32 sec_size[4];
 	u32 sec_off[4];