Work with file objects instead of bytes

When dealing with cramfs and squashfs we can use the PAK's file object
directly instead of loading the section in memory. Unfortunately UBI is
a bit harder to work with.

Author: AT0myks

reolinkfw/__init__.py

@@ -56,17 +56,6 @@ async def download(url):
             return await resp.read() if resp.status == 200 else resp.status
 
 
-def extract_fs(pak: PAK):
-    """Return the fs.bin, app.bin or rootfs.bin file as bytes."""
-    sections = {s.name: s for s in pak.sections if s.name in FS_SECTIONS}
-    if len(sections) == 2:
-        return pak.extract_section(sections["app"])
-    elif len(sections) == 1:
-        return pak.extract_section(sections.popitem()[1])
-    else:
-        return "No section found"
-
-
 def extract_paks(zip) -> list[tuple[str, PAK]]:
     """Return a list of tuples, one for each PAK file found in the ZIP.
 
@@ -94,12 +83,12 @@ def get_info_from_files(files):
     return info
 
 
-def get_files_from_squashfs(binbytes):
+def get_files_from_squashfs(fd, offset=0, closefd=True):
     # Firmwares using squashfs have either one or two file system
     # sections. When there is only one, the app directory is located at
     # /mnt/app. Otherwise it's the same as with cramfs and ubifs.
     files = dict.fromkeys(FILES)
-    with SquashFsImage.from_bytes(binbytes) as image:
+    with SquashFsImage(fd, offset, closefd) as image:
         for name in files:
             path2 = posixpath.join("/mnt/app", name)
             if (file := (image.select(name) or image.select(path2))) is not None:
@@ -120,22 +109,22 @@ def get_files_from_ubifs(binbytes):
     return files
 
 
-def get_files_from_ubi(binbytes):
-    fsbytes = get_fs_from_ubi(binbytes)
+def get_files_from_ubi(fd, size, offset=0):
+    fsbytes = get_fs_from_ubi(fd, size, offset)
     fs = FileSystem.from_magic(fsbytes[:4])
     if fs == FileSystem.UBIFS:
         return get_files_from_ubifs(fsbytes)
     elif fs == FileSystem.SQUASHFS:
-        return get_files_from_squashfs(fsbytes)
+        return get_files_from_squashfs(io.BytesIO(fsbytes))
     raise Exception("Unknown file system in UBI")
 
 
-def get_files_from_cramfs(binbytes):
+def get_files_from_cramfs(fd, offset=0, closefd=True):
     # For now all firmwares using cramfs have two file system sections.
     # 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 Cramfs.from_bytes(binbytes) as cramfs:
+    with Cramfs.from_fd(fd, offset, closefd) as cramfs:
         for name in files:
             if (file := cramfs.select(name)) is not None:
                 files[name] = file.read_bytes()
@@ -152,17 +141,18 @@ def is_local_file(string):
 
 async def get_info_from_pak(pak: PAK):
     ha = await asyncio.to_thread(sha256_pak, pak)
-    binbytes = await asyncio.to_thread(extract_fs, pak)
-    if isinstance(binbytes, str):
-        return {"error": binbytes, "sha256": ha}
-    func = {
-        FileSystem.CRAMFS: get_files_from_cramfs,
-        FileSystem.UBI: get_files_from_ubi,
-        FileSystem.SQUASHFS: get_files_from_squashfs,
-    }.get(FileSystem.from_magic(binbytes[:4]))
-    if func is None:
+    fs_sections = [s for s in pak.sections if s.name in FS_SECTIONS]
+    app = fs_sections[-1]
+    pak._fd.seek(app.start)
+    fs = FileSystem.from_magic(pak._fd.read(4))
+    if fs == FileSystem.CRAMFS:
+        files = await asyncio.to_thread(get_files_from_cramfs, pak._fd, app.start, False)
+    elif fs == FileSystem.UBI:
+        files = await asyncio.to_thread(get_files_from_ubi, pak._fd, app.len, app.start)
+    elif fs == FileSystem.SQUASHFS:
+        files = await asyncio.to_thread(get_files_from_squashfs, pak._fd, app.start, False)
+    else:
         return {"error": "Unrecognized image type", "sha256": ha}
-    files = await asyncio.to_thread(func, binbytes)
     info = get_info_from_files(files)
     return {**info, "sha256": ha}
 

reolinkfw/extract.py

@@ -2,7 +2,7 @@
 from io import StringIO
 from pathlib import Path
 
-from pakler import PAK
+from pakler import PAK, Section
 from pycramfs import Cramfs
 from pycramfs.extract import extract_dir as extract_cramfs
 from PySquashfsImage import SquashFsImage
@@ -14,22 +14,29 @@
 from reolinkfw.util import DummyLEB, get_fs_from_ubi
 
 
-def extract_file_system(fs_bytes, dest: Path = None):
+def extract_file_system(pak: PAK, section: Section, dest: Path = None):
     dest = (Path.cwd() / "reolink_fs") if dest is None else dest
     dest.mkdir(parents=True, exist_ok=True)
-    fs = FileSystem.from_magic(fs_bytes[:4])
+    pak._fd.seek(section.start)
+    fs = FileSystem.from_magic(pak._fd.read(4))
     if fs == FileSystem.UBI:
-        extract_file_system(get_fs_from_ubi(fs_bytes), dest)
-    elif fs == FileSystem.UBIFS:
-        with DummyLEB.from_bytes(fs_bytes) as leb:
-            with redirect_stdout(StringIO()):
-                # If files already exist they are not written again.
-                extract_ubifs(ubifs(leb), dest)
+        fs_bytes = get_fs_from_ubi(pak._fd, section.len, section.start)
+        fs = FileSystem.from_magic(fs_bytes[:4])
+        if fs == FileSystem.UBIFS:
+            with DummyLEB.from_bytes(fs_bytes) as leb:
+                with redirect_stdout(StringIO()):
+                    # If files already exist they are not written again.
+                    extract_ubifs(ubifs(leb), dest)
+        elif fs == FileSystem.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 == FileSystem.SQUASHFS:
-        with SquashFsImage.from_bytes(fs_bytes) as image:
+        with SquashFsImage(pak._fd, section.start, False) as image:
             extract_squashfs(image.root, dest, True)
     elif fs == FileSystem.CRAMFS:
-        with Cramfs.from_bytes(fs_bytes) as image:
+        with Cramfs.from_fd(pak._fd, section.start, False) as image:
             extract_cramfs(image.rootdir, dest, True)
     else:
         raise Exception("Unknown file system")
@@ -45,4 +52,4 @@ def extract_pak(pak: PAK, dest: Path = None, force: bool = False):
                 outpath = dest / rootfsdir / "mnt" / "app"
             else:
                 outpath = dest / rootfsdir
-            extract_file_system(pak.extract_section(section), outpath)
+            extract_file_system(pak, section, outpath)

reolinkfw/util.py

@@ -54,8 +54,10 @@ def closing_ubifile(ubifile):
         ubifile._fhandle.close()
 
 
-def get_fs_from_ubi(binbytes):
+def get_fs_from_ubi(fd, size, offset=0) -> bytes:
     """Return the first file system that sits on top of the UBI volume."""
+    fd.seek(offset)
+    binbytes = fd.read(size)
     with TempFile(binbytes) as t:
         block_size = guess_peb_size(t)
         with closing_ubifile(ubi_file(t, block_size)) as ubifile: