liblk: General refactoring

liblk/Constants.py

@@ -1,3 +1,20 @@
+#
+# This file is part of liblk (https://github.com/R0rt1z2/liblk).
+# Copyright (c) 2024 Roger Ortiz.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+
 from enum import Enum
 
 class Magic:

liblk/Exceptions.py

@@ -1,6 +1,6 @@
 #
 # This file is part of liblk (https://github.com/R0rt1z2/liblk).
-# Copyright (c) 2023 Roger Ortiz.
+# Copyright (c) 2024 Roger Ortiz.
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by

liblk/LkImage.py

@@ -1,6 +1,6 @@
 #
 # This file is part of liblk (https://github.com/R0rt1z2/liblk).
-# Copyright (c) 2023 Roger Ortiz.
+# Copyright (c) 2024 Roger Ortiz.
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -16,88 +16,63 @@
 #
 
 from pathlib import Path
+from typing import Union, List, Dict, Optional
 
 from liblk.Exceptions import InvalidLkPartition, NeedleNotFoundException
 from liblk.Constants import Magic
 from liblk.structure.LkPartition import LkPartition
 
 
 class LkImage:
-    def __init__(self, lk: str | bytes | bytearray | Path, rename_duplicates=True) -> None:
-        '''Initialize the LK image.'''
-        if isinstance(lk, str) or isinstance(lk, Path):
-            self.lk_path : str = str(lk)
-            self.lk_contents: bytearray = self.load_image(lk)
-        elif isinstance(lk, bytes):
-            self.lk_contents: bytearray = bytearray(lk)
-        elif isinstance(lk, bytearray):
-            self.lk_contents: bytearray = lk
-
-        self.partitions: list = []  # Initialize the partitions list.
-
-        # Parse the partitions.
+    def __init__(self, lk: Union[str, bytes, bytearray, Path], rename_duplicates: bool = True) -> None:
+        self.lk_path: Optional[str] = str(lk) if isinstance(lk, (str, Path)) else None
+        self.lk_contents: bytearray = self.load_image(lk) if isinstance(lk, (str, Path)) else bytearray(lk)
+        self.partitions: List[Dict[str, Union[str, LkPartition]]] = []
         self.parse_partitions(rename_duplicates)
 
-    def load_image(self, lk: str) -> bytearray:
-        '''Load the LK image.'''
-        with open(lk, "rb") as lk_file:
-            return bytearray(lk_file.read())
+    def load_image(self, lk: Union[str, Path]) -> bytearray:
+        '''Load the LK image from a file.'''
+        return bytearray(Path(lk).read_bytes())
 
-    def parse_partitions(self, rename_duplicates=False) -> None:
-        '''Parse the LK image partitions.'''
+    def parse_partitions(self, rename_duplicates: bool = False) -> None:
+        '''Parse the LK image and extract the partitions.'''
         offset = 0
-        name_counts = {}
+        name_counts: Dict[str, int] = {}
 
-        # Loop over the partition contents until we reach the end of the image, or we find a
-        # partition with a positive list_end value.
         while offset < len(self.lk_contents):
-            partition = LkPartition.from_bytes(self.lk_contents[offset:], offset)
-
-            # 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:
-                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)").
+            try:
+                partition = LkPartition.from_bytes(self.lk_contents[offset:], offset)
+            except Exception as e:
+                raise InvalidLkPartition(f"Failed to parse partition at offset 0x{offset:x}: {e}")
+
+            if partition.header.magic != Magic.MAGIC:
+                if len(self.partitions) != 0 and self.partitions[-1]["name"].lower() == "lk":
+                    break # Special case for LK images with no partition table.
+                raise InvalidLkPartition(f"Invalid magic 0x{partition.header.magic:x} at offset 0x{offset:x}")
+
             name = partition.header.name
             if rename_duplicates:
-                if name in name_counts:
-                    name_counts[name] += 1
-                    name = f"{name} ({name_counts[name]})"
-                else:
-                    name_counts[name] = 0
+                count = name_counts.get(name, 0)
+                name = f"{name} ({count})" if count else name
+                name_counts[name] = count + 1
 
-            # Once we're sure the partition is valid and the name is unique, add it to the list.
             self.partitions.append({"name": name, "partition": partition})
-
-            if partition.has_ext:
-                # The external header has a property which can be used to determine whether the
-                # current partition is the last one or not. If it is, we can stop parsing.
-                if partition.ext_header.image_list_end:
-                    break
-
             offset = partition.end_offset
 
-    def apply_patch(self, needle: str | bytes | bytearray, patch: str | bytes | bytearray) -> None:
+            if partition.has_ext and partition.ext_header.image_list_end:
+                break
+
+    def apply_patch(self, needle: Union[str, bytes, bytearray],
+                    patch: Union[str, bytes, bytearray]) -> None:
         '''Apply a binary patch to the LK image.'''
-        if isinstance(needle, str):
-            needle = bytes.fromhex(needle)
-        if isinstance(patch, str):
-            patch = bytes.fromhex(patch)
-
-        # First of all make sure the needle is actually present in the LK image. If it's not,
-        # there's no point in applying the patch.
-        offset = self.lk_contents.find(needle)
-        if offset != -1:
-            # Replace the needle with the patch. We don't need to worry about the length of the
-            # patch, since maybe the user wants to replace a 4-byte needle with a 1-byte patch.
-            self.lk_contents[offset:offset + len(patch)] = patch
-        else:
-            raise NeedleNotFoundException(needle)
+        needle_b = bytes.fromhex(needle) if isinstance(needle, str) else needle
+        patch_b = bytes.fromhex(patch) if isinstance(patch, str) else patch
+
+        offset = self.lk_contents.find(needle_b)
+        if offset == -1:
+            raise NeedleNotFoundException(needle_b)
+
+        self.lk_contents[offset:offset + len(patch_b)] = patch_b
 
     def get_partition_list(self) -> list:
         '''List all partition names.'''

liblk/__init__.py

@@ -1,6 +1,6 @@
 #
 # This file is part of liblk (https://github.com/R0rt1z2/liblk).
-# Copyright (c) 2023 Roger Ortiz.
+# Copyright (c) 2024 Roger Ortiz.
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by

liblk/structure/LkExtHeader.py

@@ -20,26 +20,17 @@
 class LkExtHeader:
     def __init__(self, header_size: int, header_version: int, image_type: int, image_list_end: int, alignment: int,
                  data_size: int, memory_address: int) -> None:
-        '''Initialize the LK extension header.'''
-        self.header_size: int = header_size
-        self.header_version: int = header_version
-        self.image_type: int = image_type
-        self.image_list_end: int = image_list_end
-        self.alignment: int = alignment
-        self.data_size: int = data_size
-        self.memory_address: int = memory_address
+        self.header_size = header_size
+        self.header_version = header_version
+        self.image_type = image_type
+        self.image_list_end = image_list_end
+        self.alignment = alignment
+        self.data_size = data_size
+        self.memory_address = memory_address
 
     @classmethod
-    def from_bytes(cls, contents: bytes):
-        '''Initialize the LK extension header from bytes.'''
-        header_size = struct.unpack_from('<I', contents)[0]
-        header_version = struct.unpack_from('<I', contents, 4)[0]
-        image_type = struct.unpack_from('<I', contents, 8)[0]
-        image_list_end = struct.unpack_from('<I', contents, 12)[0]
-        alignment = struct.unpack_from('<I', contents, 16)[0]
-        data_size = struct.unpack_from('<I', contents, 20)[0] << 32
-        memory_address = struct.unpack_from('<I', contents, 24)[0] << 32
-        return cls(header_size, header_version, image_type, image_list_end, alignment, data_size, memory_address)
+    def from_bytes(cls, contents: bytes) -> "LkExtHeader":
+        return cls(*struct.unpack_from('<IIIIIII', contents))
 
     def __str__(self) -> str:
         return (

liblk/structure/LkHeader.py

@@ -1,6 +1,6 @@
 #
 # This file is part of liblk (https://github.com/R0rt1z2/liblk).
-# Copyright (c) 2023 Roger Ortiz.
+# Copyright (c) 2024 Roger Ortiz.
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -15,40 +15,43 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-import ctypes
 import struct
 
 from liblk.Constants import Pattern
 
 class LkHeader:
     '''LK header structure.'''
     def __init__(self, magic: int, data_size: int, name: str, addressing_mode: int, memory_address: int) -> None:
-        '''Initialize the LK header structure.'''
-        self.magic: int = magic
-        self.data_size: int = data_size
-        self.name: str = name
-        self.addressing_mode: int = addressing_mode
-        self.memory_address: int = memory_address
+        self.magic = magic
+        self.data_size = data_size
+        self.name = name
+        self.addressing_mode = addressing_mode
+        self.memory_address = memory_address
 
     @classmethod
     def from_bytes(cls, contents: bytes, start_offset: int) -> "LkHeader":
         '''Create a LK header structure from bytes.'''
-        magic = struct.unpack_from('<I', contents, start_offset)[0]
-        data_size = struct.unpack_from('<I', contents, start_offset + 4)[0]
-        name = contents[start_offset + 8:start_offset + 40].decode('utf-8').rstrip('\0')
-        addressing_mode = ctypes.c_int(struct.unpack_from('<I', contents, start_offset + 40)[0]).value
-        memory_address = struct.unpack_from('<I', contents, start_offset + 44)[0]
+        unpacked_data = struct.unpack_from('<I4s32sII', contents, start_offset)
+        magic, data_size, name_bytes, addressing_mode, memory_address = unpacked_data
+        name = name_bytes.decode('utf-8').rstrip('\0')
 
-        if name == 'lk' and (memory_address & 0xffffffff) == 0xffffffff:
-            if contents.find(Pattern.LOADADDR) != -1:
-                memory_address = struct.unpack_from('<I', contents, contents.find(Pattern.LOADADDR) + 8)[0]
+        if name == 'lk' and memory_address == 0xffffffff:
+            memory_address = cls.find_secondary_memory_address(contents)
 
         return cls(magic, data_size, name, addressing_mode, memory_address)
 
+    @staticmethod
+    def find_secondary_memory_address(contents: bytes) -> int:
+        '''Adjust the memory address for 'lk' partitions if a specific pattern is found.'''
+        load_addr_position = contents.find(Pattern.LOADADDR)
+        if load_addr_position != -1:
+            return struct.unpack_from('<I', contents, load_addr_position + 8)[0]
+        return 0xffffffff
+
     def __str__(self) -> str:
         return (
             f"| {'Partition Name':<15}: {self.name}\n"
             f"| {'Data Size':<15}: {self.data_size}\n"
             f"| {'Addressing Mode':<15}: {self.addressing_mode}\n"
-            f"| {'Address':<15}: {self.memory_address:#x}\n"
+            f"| {'Memory Address':<15}: {self.memory_address:#x}\n"
         )

liblk/structure/LkPartition.py

@@ -1,6 +1,6 @@
 #
 # This file is part of liblk (https://github.com/R0rt1z2/liblk).
-# Copyright (c) 2023 Roger Ortiz.
+# Copyright (c) 2024 Roger Ortiz.
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -16,40 +16,48 @@
 #
 
 import struct
+from typing import Optional
 
 from liblk.Constants import Magic
 from liblk.structure.LkExtHeader import LkExtHeader
 from liblk.structure.LkHeader import LkHeader
 
+
 class LkPartition:
-    '''LK partition structure.'''
-    def __init__(self, header: LkHeader, has_ext: bool, data: bytearray, end_offset: int, ext_header: LkExtHeader = None) -> None:
-        '''Initialize the LK partition structure.'''
-        self.header: LkHeader = header
-        self.has_ext: bool = has_ext
-        self.data: bytearray = data
-        self.end_offset: int = end_offset
-        self.ext_header: LkExtHeader = ext_header
+    '''LK partition structure encapsulating partition data and headers.'''
+
+    def __init__(self, header: LkHeader, has_ext: bool, data: bytearray, end_offset: int,
+                 ext_header: Optional[LkExtHeader] = None) -> None:
+        self.header = header
+        self.has_ext = has_ext
+        self.data = data
+        self.end_offset = end_offset
+        self.ext_header = ext_header
 
     @classmethod
     def from_bytes(cls, contents: bytes, offset: int = 0) -> "LkPartition":
-        '''Create a LK partition structure from bytes.'''
+        '''Create a LK partition structure from bytes, handling optional extended headers.'''
+        # Determine start offset based on a specific signature in the contents
         start_offset = 0x4040 if contents[:4] == b'BFBF' else 0
+
         header = LkHeader.from_bytes(contents, start_offset)
         has_ext = struct.unpack_from('<I', contents, start_offset + 48)[0] == Magic.EXT_MAGIC
         ext_header = LkExtHeader.from_bytes(contents[start_offset + 52:]) if has_ext else None
 
         data_size = header.data_size | (ext_header.data_size << 32) if has_ext else header.data_size
         header_size = ext_header.header_size if has_ext else 512
         end_offset = offset + header_size + data_size
+
+        # Adjust end_offset for alignment if necessary
         alignment = ext_header.alignment if has_ext else 8
         if alignment and end_offset % alignment:
             end_offset += alignment - end_offset % alignment
 
-        # Ideally, if we wanted to keep the header as well (inside the data) we would simply do
-        # contents[:header_size + data_size] but since we are not keeping the header, we skip the
-        # header size and only keep the data.
-        return cls(header, has_ext, contents[header_size:header_size + data_size], end_offset, ext_header)
+        data_start = header_size
+        data_end = header_size + data_size
+        data = contents[data_start:data_end]
+
+        return cls(header, has_ext, bytearray(data), end_offset, ext_header)
 
-    def __str__(self):
-        return f"{self.header}" + (f"{self.ext_header}" if self.ext_header else "")
+    def __str__(self) -> str:
+        return f"{self.header}" + (f", {self.ext_header}" if self.has_ext else "")

liblk/structure/__init__.py

@@ -1,6 +1,6 @@
 #
 # This file is part of liblk (https://github.com/R0rt1z2/liblk).
-# Copyright (c) 2023 Roger Ortiz.
+# Copyright (c) 2024 Roger Ortiz.
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by