UBIFS changes

- Add (incomplete) API for easier interaction with UBIFS
- DummyLEB is basically the same as pre-existing ubi_file
- Using leb_virtual_file in get_fs_from_ubi was not necessary

Author: AT0myks

reolinkfw/__init__.py

@@ -13,12 +13,11 @@
 from pakler import PAK, Section, is_pak_file
 from pycramfs import Cramfs
 from PySquashfsImage import SquashFsImage
-from ubireader.ubifs import ubifs, walk
-from ubireader.ubifs.output import _process_reg_file
 
+from reolinkfw.tmpfile import TempFile
+from reolinkfw.ubifs import UBIFS
 from reolinkfw.uboot import get_arch_name, get_uboot_version, get_uimage_header
 from reolinkfw.util import (
-    DummyLEB,
     FileType,
     get_cache_file,
     get_fs_from_ubi,
@@ -87,15 +86,15 @@ def get_files_from_squashfs(fd, offset=0, closefd=True):
 
 
 def get_files_from_ubifs(binbytes):
+    # For now all firmwares using ubifs 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 DummyLEB.from_bytes(binbytes) as leb:
-        image = ubifs(leb)
-        inodes = {}
-        bad_blocks = []
-        walk.index(image, image.master_node.root_lnum, image.master_node.root_offs, inodes, bad_blocks)
-        for dent in inodes[1]['dent']:
-            if dent.name in files:
-                files[dent.name] = _process_reg_file(image, inodes[dent.inum], None)
+    with TempFile(binbytes) as file:
+        with UBIFS.from_file(file) as image:
+            for name in files:
+                if (file := image.select(name)) is not None:
+                    files[name] = file.read_bytes()
     return files
 
 

reolinkfw/extract.py

@@ -9,9 +9,12 @@
 from PySquashfsImage.extract import extract_dir as extract_squashfs
 from ubireader.ubifs import ubifs
 from ubireader.ubifs.output import extract_files as extract_ubifs
+from ubireader.ubi_io import ubi_file
+from ubireader.utils import guess_leb_size
 
 from reolinkfw import FS_SECTIONS, ROOTFS_SECTIONS
-from reolinkfw.util import DummyLEB, FileType, get_fs_from_ubi
+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: Path = None):
@@ -23,10 +26,12 @@ def extract_file_system(pak: PAK, section: Section, dest: Path = None):
         fs_bytes = get_fs_from_ubi(pak._fd, section.len, section.start)
         fs = FileType.from_magic(fs_bytes[:4])
         if fs == FileType.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)
+            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)

reolinkfw/ubifs.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from pathlib import PurePosixPath
+from stat import filemode
+from typing import Iterator, Literal, Optional
+
+from ubireader.ubifs import ubifs, walk
+from ubireader.ubifs.defines import (
+    UBIFS_COMMON_HDR_SZ,
+    UBIFS_ITYPE_BLK,
+    UBIFS_ITYPE_CHR,
+    UBIFS_ITYPE_DIR,
+    UBIFS_ITYPE_FIFO,
+    UBIFS_ITYPE_LNK,
+    UBIFS_ITYPE_REG,
+    UBIFS_ITYPE_SOCK,
+    UBIFS_NODE_MAGIC,
+    UBIFS_ROOT_INO,
+    UBIFS_SB_NODE,
+    UBIFS_SB_NODE_SZ,
+)
+from ubireader.ubifs.nodes import common_hdr, ino_node, sb_node
+from ubireader.ubifs.output import _process_reg_file
+from ubireader.ubi_io import ubi_file
+from ubireader.utils import guess_leb_size
+
+from reolinkfw.tmpfile import TempFile
+
+
+class File:
+
+    def __init__(self, image: UBIFS, nodes: dict, name: str = '', parent: Optional[Directory] = None) -> None:
+        self._image = image
+        self._nodes = nodes
+        self._name = name
+        self._parent = parent
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def parent(self) -> Optional[Directory]:
+        return self._parent
+
+    @property
+    def path(self) -> PurePosixPath:
+        if self._parent is None:
+            return PurePosixPath('/')
+        return self._parent.path / self._name
+
+    @property
+    def inode(self) -> ino_node:
+        return self._nodes["ino"]
+
+    @property
+    def mode(self) -> int:
+        return self.inode.mode
+
+    @property
+    def filemode(self) -> str:
+        return filemode(self.mode)
+
+    @property
+    def is_dir(self) -> bool:
+        return False
+
+    @property
+    def is_file(self) -> bool:
+        return False
+
+    @property
+    def is_symlink(self) -> bool:
+        return False
+
+
+class DataFile(File):
+
+    def read_bytes(self) -> bytes:
+        ...
+
+    def read_text(self, encoding: str = "utf8", errors: str = "strict") -> str:
+        return self.read_bytes().decode(encoding, errors)
+
+
+class RegularFile(DataFile):
+
+    @property
+    def is_file(self) -> Literal[True]:
+        return True
+
+    def read_bytes(self) -> bytes:
+        return _process_reg_file(self._image._ubifs, self._nodes, None)
+
+
+class Symlink(DataFile):
+
+    @property
+    def is_symlink(self) -> Literal[True]:
+        return True
+
+    def read_bytes(self) -> bytes:
+        return self.inode.data
+
+    def readlink(self) -> PurePosixPath:
+        return PurePosixPath(self.read_text())
+
+
+class Directory(File):
+
+    def __init__(self, image: UBIFS, nodes: dict, name: str = '', parent: Optional[Directory] = None) -> None:
+        super().__init__(image, nodes, name, parent)
+        self._children = {}
+        for dent in nodes.get("dent", []):
+            cls = filetype[dent.type]
+            self._children[dent.name] = cls(image, image.inodes[dent.inum], dent.name, self)
+
+    def __iter__(self) -> Iterator[File]:
+        yield from self._children.values()
+
+    @property
+    def is_dir(self) -> Literal[True]:
+        return True
+
+    @property
+    def children(self) -> dict[str, File]:
+        return self._children
+
+    def select(self, path) -> Optional[File]:
+        """Select a file of any kind by path.
+
+        The path can be absolute or relative.
+        Special entries `'.'` and `'..'` are supported.
+        """
+        path = PurePosixPath(path)
+        if str(path) == "..":
+            return self.parent if self.parent is not None else self
+        if path.root == '/':
+            if str(self.path) == '/':
+                path = path.relative_to('/')
+            else:
+                return self._image.root.select(path)
+        if str(path) == '.':
+            return self
+        child, *descendants = path.parts
+        if (file := self._children.get(child)) is not None:
+            if isinstance(file, Directory) and descendants:
+                return file.select(PurePosixPath(*descendants))
+            elif not descendants:
+                return file
+        return None
+
+    def riter(self) -> Iterator[File]:
+        yield self
+        for file in self:
+            if isinstance(file, Directory):
+                yield from file.riter()
+            else:
+                yield file
+
+
+class UBIFS:
+
+    def __init__(self, ubifs: ubifs) -> None:
+        self._ubifs = ubifs
+        self._inodes = {}
+        self._bad_blocks = []
+        walk.index(ubifs, ubifs.master_node.root_lnum, ubifs.master_node.root_offs, self._inodes, self._bad_blocks)
+        self._root = Directory(self, self._inodes[UBIFS_ROOT_INO])
+
+    def __enter__(self) -> UBIFS:
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.close()
+
+    def __iter__(self) -> Iterator[File]:
+        yield from self._root.riter()
+
+    @property
+    def inodes(self) -> dict:
+        return self._inodes
+
+    @property
+    def root(self) -> Directory:
+        return self._root
+
+    def close(self) -> None:
+        self._ubifs._file._fhandle.close()
+
+    def select(self, path) -> Optional[File]:
+        return self._root.select(path)
+
+    @classmethod
+    def from_file(cls, path) -> UBIFS:
+        return cls(ubifs(ubi_file(path, guess_leb_size(path))))
+
+    @classmethod
+    def from_bytes(cls, bytes_, offset: int = 0) -> UBIFS:
+        chdr = common_hdr(bytes_[offset:offset+UBIFS_COMMON_HDR_SZ])
+        if chdr.magic != int.from_bytes(UBIFS_NODE_MAGIC, "little") or chdr.node_type != UBIFS_SB_NODE:
+            raise Exception("Not UBIFS")
+        sb_start = offset + UBIFS_COMMON_HDR_SZ
+        sb_end = sb_start + UBIFS_SB_NODE_SZ
+        sblk = sb_node(bytes_[sb_start:sb_end])
+        tmpfile = TempFile(bytes_[offset:])
+        tmpfile.open()  # tmpfile's close() will not be called.
+        return cls(ubifs(ubi_file(tmpfile, sblk.leb_size)))
+
+
+filetype = {
+    UBIFS_ITYPE_REG: RegularFile,
+    UBIFS_ITYPE_DIR: Directory,
+    UBIFS_ITYPE_LNK: Symlink,
+    UBIFS_ITYPE_BLK: File,
+    UBIFS_ITYPE_CHR: File,
+    UBIFS_ITYPE_FIFO: File,
+    UBIFS_ITYPE_SOCK: File,
+}

reolinkfw/util.py

@@ -1,5 +1,4 @@
 import hashlib
-import io
 from contextlib import contextmanager
 from enum import Enum
 from functools import partial
@@ -14,7 +13,7 @@
 from PySquashfsImage.const import SQUASHFS_MAGIC
 from ubireader.ubi import ubi
 from ubireader.ubi.defines import UBI_EC_HDR_MAGIC as UBI_MAGIC
-from ubireader.ubi_io import ubi_file, leb_virtual_file
+from ubireader.ubi_io import ubi_file
 from ubireader.ubifs.defines import UBIFS_NODE_MAGIC as UBIFS_MAGIC
 from ubireader.utils import guess_peb_size
 
@@ -39,41 +38,6 @@ def from_magic(cls, key, default=None):
             return default
 
 
-class DummyLEB:
-    """A class that emulates ubireader's `leb_virtual_file`."""
-
-    def __init__(self, fd):
-        self._fd = fd
-        self._last_read_addr = 0
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.close()
-
-    def close(self):
-        self._fd.close()
-
-    def read(self, size):
-        self._last_read_addr = self._fd.tell()
-        return self._fd.read(size)
-
-    def reset(self):
-        return self._fd.seek(0)
-
-    def seek(self, offset):
-        return self._fd.seek(offset)
-
-    def last_read_addr(self):
-        """Start address of last physical file read."""
-        return self._last_read_addr
-
-    @classmethod
-    def from_bytes(cls, bytes_):
-        return cls(io.BytesIO(bytes_))
-
-
 @contextmanager
 def closing_ubifile(ubifile):
     try:
@@ -90,8 +54,8 @@ def get_fs_from_ubi(fd, size, offset=0) -> bytes:
         block_size = guess_peb_size(t)
         with closing_ubifile(ubi_file(t, block_size)) as ubifile:
             ubi_obj = ubi(ubifile)
-            vol_blocks = ubi_obj.images[0].volumes.popitem()[1].get_blocks(ubi_obj.blocks)
-            return b''.join(leb_virtual_file(ubi_obj, vol_blocks).reader())
+            volume = ubi_obj.images[0].volumes.popitem()[1]
+            return b''.join(volume.reader(ubi_obj))
 
 
 def sha256_pak(pak: PAK) -> str: