New ReolinkFirmware class

Author: AT0myks

README.md

@@ -154,16 +154,15 @@ You should not use it to repack a custom firmware.
 ### As a library
 
 ```py
-from pakler import PAK
-from reolinkfw import get_info
+from reolinkfw import ReolinkFirmware, get_info
 from reolinkfw.extract import extract_pak
 
 url = "https://reolink-storage.s3.amazonaws.com/website/firmware/20200523firmware/RLC-410-5MP_20_20052300.zip"
 print(get_info(url))
 file = "/home/ben/RLC-410-5MP_20_20052300.zip"
 print(get_info(file))
-with PAK.from_file(file) as pak:
-    extract_pak(pak)
+with ReolinkFirmware.from_file(file) as fw:
+    extract_pak(fw)
 ```
 
 In most cases where a URL is used, it will be a direct link to the file

reolinkfw/__init__.py

@@ -2,7 +2,7 @@
 import io
 import posixpath
 import re
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Iterator, Mapping
 from pathlib import Path
 from typing import IO, Any, BinaryIO, Optional, Union
 from urllib.parse import parse_qsl, urlparse
@@ -22,6 +22,7 @@
 from reolinkfw.uboot import get_arch_name, get_uboot_version, get_uimage_header
 from reolinkfw.util import (
     FileType,
+    SectionFile,
     get_cache_file,
     get_fs_from_ubi,
     has_cache,
@@ -34,8 +35,63 @@
 FILES = ("version_file", "version.json", "dvr.xml", "dvr", "router")
 INFO_KEYS = ("firmware_version_prefix", "board_type", "board_name", "build_date", "display_type_info", "detail_machine_type", "type")
 
-ROOTFS_SECTIONS = ["fs", "rootfs"]
-FS_SECTIONS = ROOTFS_SECTIONS + ["app"]
+UBOOT_SECTIONS = ("uboot", "uboot1", "BOOT")
+KERNEL_SECTIONS = ("kernel", "KERNEL")
+ROOTFS_SECTIONS = ("fs", "rootfs")
+FS_SECTIONS = ROOTFS_SECTIONS + ("app",)
+
+
+class ReolinkFirmware(PAK):
+
+    def __init__(self, fd: BinaryIO, offset: int = 0, closefd: bool = True) -> None:
+        super().__init__(fd, offset, closefd)
+        self._uboot_section_name = self._get_uboot_section_name()
+        self._kernel_section_name = self._get_kernel_section_name()
+        self._sdict = {s.name: s for s in self}
+        self._open_files = 1
+
+    def __del__(self) -> None:
+        self.close()
+
+    def __getitem__(self, key: Union[int, str]) -> Section:
+        if isinstance(key, int):
+            return self.sections[key]
+        if isinstance(key, str):
+            if key.lower() == "uboot":
+                key = self._uboot_section_name
+            elif key.lower() == "kernel":
+                key = self._kernel_section_name
+            return self._sdict[key]
+        raise TypeError
+
+    def __iter__(self) -> Iterator[Section]:
+        yield from self.sections
+
+    def _fdclose(self, fd: BinaryIO) -> None:
+        self._open_files -= 1
+        if self._closefd and not self._open_files:
+            fd.close()
+
+    def _get_uboot_section_name(self) -> str:
+        for section in self:
+            if section.len and section.name in UBOOT_SECTIONS:
+                return section.name
+        raise Exception("U-Boot section not found")
+
+    def _get_kernel_section_name(self) -> str:
+        for section in self:
+            if section.len and section.name in KERNEL_SECTIONS:
+                return section.name
+        raise Exception("Kernel section not found")
+
+    def open(self, section: Section) -> SectionFile:
+        self._open_files += 1
+        return SectionFile(self._fd, section, self._fdclose)
+
+    def close(self) -> None:
+        if self._fd is not None:
+            self._fdclose(self._fd)
+            self._fd = None
 
 
 async def download(url: StrOrURL) -> Union[bytes, int]:
@@ -48,7 +104,7 @@ async def download(url: StrOrURL) -> Union[bytes, int]:
             return await resp.read() if resp.status == 200 else resp.status
 
 
-def extract_paks(zip: Union[StrPath, IO[bytes]]) -> list[tuple[str, PAK]]:
+def extract_paks(zip: Union[StrPath, IO[bytes]]) -> list[tuple[str, ReolinkFirmware]]:
     """Return a list of tuples, one for each PAK file found in the ZIP.
 
     It is the caller's responsibility to close the PAK files.
@@ -58,7 +114,7 @@ def extract_paks(zip: Union[StrPath, IO[bytes]]) -> list[tuple[str, PAK]]:
         for name in myzip.namelist():
             file = myzip.open(name)
             if is_pak_file(file):
-                paks.append((file.name, PAK.from_fd(file)))
+                paks.append((file.name, ReolinkFirmware.from_fd(file)))
             else:
                 file.close()
     return paks
@@ -93,8 +149,8 @@ def get_files_from_ubifs(binbytes: Buffer) -> dict[Files, Optional[bytes]]:
     # The interesting files are in the root directory of the "app" one.
     # Using select() with a relative path is enough.
     files = dict.fromkeys(FILES)
-    with TempFile(binbytes) as file:
-        with UBIFS.from_file(file) as image:
+    with TempFile(binbytes) as tempfile:
+        with UBIFS.from_file(tempfile) as image:
             for name in files:
                 if (file := image.select(name)) is not None:
                     files[name] = file.read_bytes()
@@ -131,43 +187,43 @@ def is_local_file(string: StrPath) -> bool:
     return Path(string).is_file()
 
 
-def get_fs_info(pak: PAK, fs_sections: Iterable[Section]) -> list[dict[str, str]]:
+def get_fs_info(fw: ReolinkFirmware, fs_sections: Iterable[Section]) -> list[dict[str, str]]:
     result = []
     for section in fs_sections:
-        pak._fd.seek(section.start)
-        fs = FileType.from_magic(pak._fd.read(4))
-        if fs == FileType.UBI:
-            pak._fd.seek(section.start + 266240)
-            fs = FileType.from_magic(pak._fd.read(4))
+        with fw.open(section) as f:
+            fs = FileType.from_magic(f.read(4))
+            if fs == FileType.UBI:
+                f.seek(266240)
+                fs = FileType.from_magic(f.read(4))
         result.append({
             "name": section.name,
             "type": fs.name.lower() if fs is not None else "unknown"
         })
     return result
 
 
-async def get_info_from_pak(pak: PAK) -> dict[str, Any]:
-    ha = await asyncio.to_thread(sha256_pak, pak)
-    fs_sections = [s for s in pak.sections if s.name in FS_SECTIONS]
+async def get_info_from_pak(fw: ReolinkFirmware) -> dict[str, Any]:
+    ha = await asyncio.to_thread(sha256_pak, fw)
+    fs_sections = [s for s in fw if s.name in FS_SECTIONS]
     app = fs_sections[-1]
-    pak._fd.seek(app.start)
-    fs = FileType.from_magic(pak._fd.read(4))
-    if fs == FileType.CRAMFS:
-        files = await asyncio.to_thread(get_files_from_cramfs, pak._fd, app.start, False)
-    elif fs == FileType.UBI:
-        files = await asyncio.to_thread(get_files_from_ubi, pak._fd, app.len, app.start)
-    elif fs == FileType.SQUASHFS:
-        files = await asyncio.to_thread(get_files_from_squashfs, pak._fd, app.start, False)
-    else:
-        return {"error": "Unrecognized image type", "sha256": ha}
-    uimage = get_uimage_header(pak)
+    with fw.open(app) as f:
+        fs = FileType.from_magic(f.read(4))
+        if fs == FileType.CRAMFS:
+            files = await asyncio.to_thread(get_files_from_cramfs, f, 0, False)
+        elif fs == FileType.UBI:
+            files = await asyncio.to_thread(get_files_from_ubi, f, app.len, 0)
+        elif fs == FileType.SQUASHFS:
+            files = await asyncio.to_thread(get_files_from_squashfs, f, 0, False)
+        else:
+            return {"error": "Unrecognized image type", "sha256": ha}
+    uimage = get_uimage_header(fw)
     return {
         **get_info_from_files(files),
         "os": "Linux" if uimage.os == 5 else "Unknown",
         "architecture": get_arch_name(uimage.arch),
         "kernel_image_name": uimage.name,
-        "uboot_version": get_uboot_version(pak),
-        "filesystems": get_fs_info(pak, fs_sections),
+        "uboot_version": get_uboot_version(fw),
+        "filesystems": get_fs_info(fw, fs_sections),
         "sha256": ha
     }
 
@@ -185,7 +241,7 @@ async def direct_download_url(url: str) -> str:
     return url
 
 
-async def get_paks(file_or_url: StrPathURL, use_cache: bool = True) -> list[tuple[Optional[str], PAK]]:
+async def get_paks(file_or_url: StrPathURL, use_cache: bool = True) -> list[tuple[Optional[str], ReolinkFirmware]]:
     """Return PAK files read from an on-disk file or a URL.
 
     The file or resource may be a ZIP or a PAK. On success return a
@@ -206,7 +262,7 @@ async def get_paks(file_or_url: StrPathURL, use_cache: bool = True) -> list[tupl
         if use_cache:
             make_cache_file(file_or_url, zip_or_pak_bytes, pakname)
         if is_pak_file(zip_or_pak_bytes):
-            return [(pakname, PAK.from_bytes(zip_or_pak_bytes))]
+            return [(pakname, ReolinkFirmware.from_bytes(zip_or_pak_bytes))]
         else:
             zipfile = io.BytesIO(zip_or_pak_bytes)
             if is_zipfile(zipfile):
@@ -218,7 +274,7 @@ async def get_paks(file_or_url: StrPathURL, use_cache: bool = True) -> list[tupl
         if is_zipfile(file_or_url):
             return await asyncio.to_thread(extract_paks, file_or_url)
         elif is_pak_file(file_or_url):
-            return [(file_or_url.name, PAK.from_file(file_or_url))]
+            return [(file_or_url.name, ReolinkFirmware.from_file(file_or_url))]
         raise Exception("Not a ZIP or a PAK file")
     raise Exception("Not a URL or file")
 

reolinkfw/extract.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import Optional
 
-from pakler import PAK, Section
+from pakler import Section
 from pycramfs import Cramfs
 from pycramfs.extract import extract_dir as extract_cramfs
 from PySquashfsImage import SquashFsImage
@@ -13,49 +13,49 @@
 from ubireader.ubi_io import ubi_file
 from ubireader.utils import guess_leb_size
 
-from reolinkfw import FS_SECTIONS, ROOTFS_SECTIONS
+from reolinkfw import FS_SECTIONS, ROOTFS_SECTIONS, ReolinkFirmware
 from reolinkfw.tmpfile import TempFile
 from reolinkfw.util import FileType, closing_ubifile, get_fs_from_ubi
 
 
-def extract_file_system(pak: PAK, section: Section, dest: Optional[Path] = None) -> None:
+def extract_file_system(fw: ReolinkFirmware, section: Section, dest: Optional[Path] = None) -> None:
     dest = (Path.cwd() / "reolink_fs") if dest is None else dest
     dest.mkdir(parents=True, exist_ok=True)
-    pak._fd.seek(section.start)
-    fs = FileType.from_magic(pak._fd.read(4))
-    if fs == FileType.UBI:
-        fs_bytes = get_fs_from_ubi(pak._fd, section.len, section.start)
-        fs = FileType.from_magic(fs_bytes[:4])
-        if fs == FileType.UBIFS:
-            with TempFile(fs_bytes) as file:
-                block_size = guess_leb_size(file)
-                with closing_ubifile(ubi_file(file, block_size)) as ubifile:
-                    with redirect_stdout(StringIO()):
-                        # Files that already exist are not written again.
-                        extract_ubifs(ubifs(ubifile), dest)
+    with fw.open(section) as f:
+        fs = FileType.from_magic(f.read(4))
+        if fs == FileType.UBI:
+            fs_bytes = get_fs_from_ubi(f, section.len, 0)
+            fs = FileType.from_magic(fs_bytes[:4])
+            if fs == FileType.UBIFS:
+                with TempFile(fs_bytes) as file:
+                    block_size = guess_leb_size(file)
+                    with closing_ubifile(ubi_file(file, block_size)) as ubifile:
+                        with redirect_stdout(StringIO()):
+                            # Files that already exist are not written again.
+                            extract_ubifs(ubifs(ubifile), dest)
+            elif fs == FileType.SQUASHFS:
+                with SquashFsImage.from_bytes(fs_bytes) as image:
+                    extract_squashfs(image.root, dest, True)
+            else:
+                raise Exception("Unknown file system in UBI")
         elif fs == FileType.SQUASHFS:
-            with SquashFsImage.from_bytes(fs_bytes) as image:
+            with SquashFsImage(f, 0, False) as image:
                 extract_squashfs(image.root, dest, True)
+        elif fs == FileType.CRAMFS:
+            with Cramfs.from_fd(f, 0, False) as image:
+                extract_cramfs(image.rootdir, dest, True)
         else:
-            raise Exception("Unknown file system in UBI")
-    elif fs == FileType.SQUASHFS:
-        with SquashFsImage(pak._fd, section.start, False) as image:
-            extract_squashfs(image.root, dest, True)
-    elif fs == FileType.CRAMFS:
-        with Cramfs.from_fd(pak._fd, section.start, False) as image:
-            extract_cramfs(image.rootdir, dest, True)
-    else:
-        raise Exception("Unknown file system")
+            raise Exception("Unknown file system")
 
 
-def extract_pak(pak: PAK, dest: Optional[Path] = None, force: bool = False) -> None:
+def extract_pak(fw: ReolinkFirmware, dest: Optional[Path] = None, force: bool = False) -> None:
     dest = (Path.cwd() / "reolink_firmware") if dest is None else dest
     dest.mkdir(parents=True, exist_ok=force)
-    rootfsdir = [s.name for s in pak.sections if s.name in ROOTFS_SECTIONS][0]
-    for section in pak.sections:
+    rootfsdir = [s.name for s in fw if s.name in ROOTFS_SECTIONS][0]
+    for section in fw:
         if section.name in FS_SECTIONS:
             if section.name == "app":
                 outpath = dest / rootfsdir / "mnt" / "app"
             else:
                 outpath = dest / rootfsdir
-            extract_file_system(pak, section, outpath)
+            extract_file_system(fw, section, outpath)

reolinkfw/uboot.py

@@ -3,11 +3,12 @@
 import re
 from ctypes import BigEndianStructure, c_char, c_uint32, c_uint8, sizeof
 from enum import IntEnum
-from typing import BinaryIO, Optional
+from typing import TYPE_CHECKING, BinaryIO, Optional
 
 import pybcl
-from pakler import PAK
 
+if TYPE_CHECKING:
+    from reolinkfw import ReolinkFirmware
 from reolinkfw.util import FileType
 
 UBOOT_MAGIC = 0x27051956
@@ -101,38 +102,30 @@ def from_fd(cls, fd: BinaryIO) -> LegacyImageHeader:
         return cls.from_buffer_copy(fd.read(sizeof(cls)))
 
 
-def is_bcl_compressed(fd: BinaryIO) -> bool:
-    size = len(pybcl.BCL_MAGIC_BYTES)
-    magic = fd.read(size)
-    fd.seek(-size, 1)
-    return magic == pybcl.BCL_MAGIC_BYTES
-
-
-def get_uboot_version(pak: PAK) -> Optional[str]:
-    for section in pak.sections:
+def get_uboot_version(fw: ReolinkFirmware) -> Optional[str]:
+    for section in fw:
         if section.len and "uboot" in section.name.lower():
             # This section is always named 'uboot' or 'uboot1'.
-            pak._fd.seek(section.start)
-            if is_bcl_compressed(pak._fd):
-                # Sometimes section.len - sizeof(hdr) is 1 to 3 bytes larger
-                # than hdr.size. The extra bytes are 0xff (padding?). This
-                # could explain why the compressed size is added to the header.
-                hdr = pybcl.HeaderVariant.from_fd(pak._fd)
-                data = pybcl.decompress(pak._fd.read(hdr.size), hdr.algo, hdr.outsize)
-            else:
-                data = pak._fd.read(section.len)
+            with fw.open(section) as f:
+                if f.peek(len(pybcl.BCL_MAGIC_BYTES)) == pybcl.BCL_MAGIC_BYTES:
+                    # Sometimes section.len - sizeof(hdr) is 1 to 3 bytes larger
+                    # than hdr.size. The extra bytes are 0xff (padding?). This
+                    # could explain why the compressed size is added to the header.
+                    hdr = pybcl.HeaderVariant.from_fd(f)
+                    data = pybcl.decompress(f.read(hdr.size), hdr.algo, hdr.outsize)
+                else:
+                    data = f.read(section.len)
             match = re.search(b"U-Boot [0-9]{4}\.[0-9]{2}.*? \(.*?\)", data)
             return match.group().decode() if match is not None else None
     return None
 
 
-def get_uimage_header(pak: PAK) -> LegacyImageHeader:
-    for section in pak.sections:
-        pak._fd.seek(section.start)
-        if section.len and FileType.from_magic(pak._fd.read(4)) == FileType.UIMAGE:
-            # This section is always named 'KERNEL' or 'kernel'.
-            pak._fd.seek(section.start)
-            return LegacyImageHeader.from_fd(pak._fd)
+def get_uimage_header(fw: ReolinkFirmware) -> LegacyImageHeader:
+    for section in fw:
+        with fw.open(section) as f:
+            if section.len and FileType.from_magic(f.peek(4)) == FileType.UIMAGE:
+                # This section is always named 'KERNEL' or 'kernel'.
+                return LegacyImageHeader.from_fd(f)
     raise Exception("No kernel section found")