From 6ba11141f610a5068c26b779aaac552d55182847 Mon Sep 17 00:00:00 2001
From: R0rt1z2 <me@r0rt1z2.com>
Date: Thu, 2 Nov 2023 12:58:16 +0100
Subject: [PATCH] liblk: Fix partition parsing for legacy LK images

* In legacy LK images, typically only a single, primary partition
  exists, identified as 'lk'. These images don't feature additional
  partitions with external headers, making it hard to guess when we
  should stop the parsing loop.

* To address this limitation, the parsing process should include a
  check to ascertain if the most recently extracted partition is 'lk'.
  This way, if we hit an invalid magic, we'll break out of the loop if
  the current LK is a legacy image.
---
 liblk/LkImage.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/liblk/LkImage.py b/liblk/LkImage.py
index c4d8d30..d4962c5 100644
--- a/liblk/LkImage.py
+++ b/liblk/LkImage.py
@@ -56,7 +56,10 @@ def parse_partitions(self, rename_duplicates=False) -> None:
             # Always make sure the partition header magic is valid, this is the only way to tell
             # whether the partition is valid or not.
             if not partition.header.magic == Magic.MAGIC:
-                raise InvalidLkPartition(f"Invalid magic 0x{partition.header.magic:x} at offset 0x{offset:x}")
+                if len(self.partitions) != 0 and self.partitions[-1]["name"] == 'lk':
+                    break # In legacy LK images, this is completely fine, so just break the loop.
+                else:
+                    raise InvalidLkPartition(f"Invalid magic 0x{partition.header.magic:x} at offset 0x{offset:x}")
 
             # There are certain cases where one partition is repeated more than once. In order
             # to avoid name collisions, append a number to the name (e.g. "cert1" -> "cert1 (1)").