feat: add type hints to ubireader.ubifs

Author: scootermon

ubireader/ubifs/__init__.py

@@ -17,10 +17,14 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
+from typing import TYPE_CHECKING
 from ubireader.debug import error, log, verbose_display
 from ubireader.ubifs.defines import *
 from ubireader.ubifs import nodes, display
-from typing import Optional
+
+if TYPE_CHECKING:
+    from ubireader.ubi_io import ubi_file as UbiFile, leb_virtual_file as LebVirtualFile
 
 class ubifs():
     """UBIFS object
@@ -36,7 +40,7 @@ class ubifs():
     Obj:mst_node       -- Master Node of UBIFS image LEB1
     Obj:mst_node2      -- Master Node 2 of UBIFS image LEB2
     """
-    def __init__(self, ubifs_file, master_key: Optional[bytes] = None):
+    def __init__(self, ubifs_file: UbiFile | LebVirtualFile, master_key: bytes | None = None) -> None:
         self.__name__ = 'UBIFS'
         self._file = ubifs_file
         self.master_key = master_key
@@ -59,7 +63,7 @@ def __init__(self, ubifs_file, master_key: Optional[bytes] = None):
         except Exception as e:
             error(self, 'Fatal', 'Super block error: %s' % e)
 
-        self._mst_nodes = [None, None]
+        self._mst_nodes: list[nodes.mst_node | None] = [None, None]
         for i in range(0, 2):
             try:
                 mst_offset = self.leb_size * (UBIFS_MST_LNUM + i) 
@@ -88,12 +92,12 @@ def __init__(self, ubifs_file, master_key: Optional[bytes] = None):
             log(self , 'Swapping Master Nodes due to bad first node.')
 
 
-    def _get_file(self):
+    def _get_file(self) -> UbiFile | LebVirtualFile:
         return self._file
     file = property(_get_file)
 
 
-    def _get_superblock(self):
+    def _get_superblock(self) -> nodes.sb_node:
         """ Superblock Node Object
 
         Returns:
@@ -103,7 +107,7 @@ def _get_superblock(self):
     superblock_node = property(_get_superblock)
 
 
-    def _get_master_node(self):
+    def _get_master_node(self) -> nodes.mst_node | None:
         """Master Node Object
 
         Returns:
@@ -113,7 +117,7 @@ def _get_master_node(self):
     master_node = property(_get_master_node)
 
 
-    def _get_master_node2(self):
+    def _get_master_node2(self) -> nodes.mst_node | None:
         """Master Node Object 2
 
         Returns:
@@ -123,7 +127,7 @@ def _get_master_node2(self):
     master_node2 = property(_get_master_node2)
 
 
-    def _get_leb_size(self):
+    def _get_leb_size(self) -> int:
         """LEB size of UBI blocks in file.
 
         Returns:
@@ -133,7 +137,7 @@ def _get_leb_size(self):
     leb_size = property(_get_leb_size)
 
 
-    def _get_min_io_size(self):
+    def _get_min_io_size(self) -> int:
         """Min I/O Size
 
         Returns:
@@ -142,7 +146,7 @@ def _get_min_io_size(self):
         return self._min_io_size
     min_io_size = property(_get_min_io_size)
 
-    def display(self, tab=''):
+    def display(self, tab: str = '') -> str:
         """Print information about this object.
         
         Argument:

ubireader/ubifs/decrypt.py

@@ -1,12 +1,19 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
 from ubireader.ubifs.defines import UBIFS_XATTR_NAME_ENCRYPTION_CONTEXT
 from ubireader.debug import error
 from cryptography.hazmat.primitives.ciphers import (
    Cipher, algorithms, modes
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from ubireader.ubifs import ubifs as Ubifs, nodes
+    from ubireader.ubifs.walk import Inode
+
 AES_BLOCK_SIZE = algorithms.AES.block_size // 8
 
-def lookup_inode_nonce(inodes: dict, inode: dict) -> bytes:
+def lookup_inode_nonce(inodes: Mapping[int, Inode], inode: Inode) -> bytes | None:
     # get the extended attribute 'xent' of the inode
     if 'xent' not in inode or not inode['xent']:
         raise ValueError(f"No xent found for inode {inode}")
@@ -29,7 +36,7 @@ def derive_key_from_nonce(master_key: bytes, nonce: bytes) -> bytes:
     return derived_key
 
 
-def filename_decrypt(key: bytes, ciphertext: bytes):
+def filename_decrypt(key: bytes, ciphertext: bytes) -> bytes:
 
     # using AES CTS-CBC mode not supported by pyca cryptography 
     if len(ciphertext) > AES_BLOCK_SIZE:
@@ -59,15 +66,15 @@ def filename_decrypt(key: bytes, ciphertext: bytes):
     return plaintext.rstrip(b'\x00')
 
 
-def datablock_decrypt(block_key: bytes, block_iv: bytes, block_data: bytes):
+def datablock_decrypt(block_key: bytes, block_iv: bytes, block_data: bytes) -> bytes:
     decryptor = Cipher(
         algorithms.AES(block_key),
         modes.XTS(block_iv),
     ).decryptor()
     return decryptor.update(block_data) + decryptor.finalize()
 
 
-def decrypt_filenames(ubifs, inodes):
+def decrypt_filenames(ubifs: Ubifs, inodes: Mapping[int, Inode]) -> None:
     if ubifs.master_key is None:
         for inode in inodes.values():
             for dent in inode.get('dent', []):
@@ -87,7 +94,7 @@ def decrypt_filenames(ubifs, inodes):
         error(decrypt_filenames, 'Error', str(e))
 
 
-def decrypt_symlink_target(ubifs, inodes, dent_node) -> str:
+def decrypt_symlink_target(ubifs: Ubifs, inodes: Mapping[int, Inode], dent_node: nodes.dent_node) -> str:
     if ubifs.master_key is None:
         return inodes[dent_node.inum]['ino'].data.decode()
     inode = inodes[dent_node.inum]

ubireader/ubifs/display.py

@@ -17,16 +17,22 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
+from typing import TYPE_CHECKING
 from ubireader.ubifs.defines import PRINT_UBIFS_FLGS, PRINT_UBIFS_MST
 
-def ubifs(ubifs, tab=''):
+if TYPE_CHECKING:
+    from ubireader.ubifs import ubifs as Ubifs
+    from ubireader.ubifs import nodes
+
+def ubifs(ubifs: Ubifs, tab: str = '') -> str:
     buf = '%sUBIFS Image\n' % (tab)
     buf += '%s---------------------\n' % (tab)
     buf += '%sMin I/O: %s\n' % (tab, ubifs.min_io_size)
     buf += '%sLEB Size: %s\n' % (tab, ubifs.leb_size)
     return buf
 
-def common_hdr(chdr, tab=''):
+def common_hdr(chdr: nodes.common_hdr, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, chdr)
     buf += '%s---------------------\n' % (tab)
     tab += '\t'
@@ -41,7 +47,7 @@ def common_hdr(chdr, tab=''):
             buf += '%s%s: %r\n' % (tab, key, value)
     return buf
 
-def sb_node(node, tab=''):
+def sb_node(node: nodes.sb_node, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, node)
     buf += '%sFile offset: %s\n' % (tab, node.file_offset)
     buf += '%s---------------------\n' % (tab)
@@ -69,7 +75,7 @@ def sb_node(node, tab=''):
     return buf
 
 
-def mst_node(node, tab=''):
+def mst_node(node: nodes.mst_node, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, node)
     buf += '%sFile offset: %s\n' % (tab, node.file_offset)
     buf += '%s---------------------\n' % (tab)
@@ -94,7 +100,7 @@ def mst_node(node, tab=''):
     return buf
 
 
-def dent_node(node, tab=''):
+def dent_node(node: nodes.dent_node, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, node)
     buf += '%s---------------------\n' % (tab)
     tab += '\t'
@@ -108,7 +114,7 @@ def dent_node(node, tab=''):
     return buf
 
 
-def data_node(node, tab=''):
+def data_node(node: nodes.data_node, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, node)
     buf += '%s---------------------\n' % (tab)
     tab += '\t'
@@ -122,7 +128,7 @@ def data_node(node, tab=''):
     return buf
 
 
-def idx_node(node, tab=''):
+def idx_node(node: nodes.idx_node, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, node)
     buf += '%s---------------------\n' % (tab)
     tab += '\t'
@@ -136,7 +142,7 @@ def idx_node(node, tab=''):
     return buf
 
 
-def ino_node(node, tab=''):
+def ino_node(node: nodes.ino_node, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, node)
     buf += '%s---------------------\n' % (tab)
     tab += '\t'
@@ -150,7 +156,7 @@ def ino_node(node, tab=''):
     return buf
 
 
-def branch(node, tab=''):
+def branch(node: nodes.branch, tab: str = '') -> str:
     buf = '%s%s\n' % (tab, node)
     buf += '%s---------------------\n' % (tab)
     tab += '\t'

ubireader/ubifs/list.py

@@ -17,25 +17,30 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
 import os
 import time
-import struct
+from typing import TYPE_CHECKING
 from ubireader.ubifs.decrypt import decrypt_symlink_target
 from ubireader.ubifs.defines import *
 from ubireader.ubifs import walk
 from ubireader.ubifs.misc import process_reg_file
 from ubireader.debug import error
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from ubireader.ubifs import ubifs as Ubifs, nodes
+    from ubireader.ubifs.walk import Inode
 
-def list_files(ubifs, list_path):
+def list_files(ubifs: Ubifs, list_path: str) -> None:
     pathnames = list_path.split("/")
-    pnames = []
+    pnames: list[str] = []
     for i in pathnames:
         if len(i) > 0:
             pnames.append(i)
     try:
-        inodes = {}
-        bad_blocks = []
+        inodes: dict[int, Inode] = {}
+        bad_blocks: list[int] = []
 
         walk.index(ubifs, ubifs.master_node.root_lnum, ubifs.master_node.root_offs, inodes, bad_blocks)
 
@@ -60,18 +65,18 @@ def list_files(ubifs, list_path):
         error(list_files, 'Error', '%s' % e)
 
 
-def copy_file(ubifs, filepath, destpath):
+def copy_file(ubifs: Ubifs, filepath: str, destpath: str) -> bool:
     pathnames = filepath.split("/")
-    pnames = []
+    pnames: list[str] = []
     for i in pathnames:
         if len(i) > 0:
             pnames.append(i)
 
     filename = pnames[len(pnames)-1]
     del pnames[-1]
 
-    inodes = {}
-    bad_blocks = []
+    inodes: dict[int, Inode] = {}
+    bad_blocks: list[int] = []
 
     walk.index(ubifs, ubifs.master_node.root_lnum, ubifs.master_node.root_offs, inodes, bad_blocks)
 
@@ -97,7 +102,7 @@ def copy_file(ubifs, filepath, destpath):
     return False
 
 
-def find_dir(inodes, inum, names, idx):
+def find_dir(inodes: Mapping[int, Inode], inum: int, names: list[str], idx: int) -> int | None:
     if len(names) == 0:
         return 1
     for dent in inodes[inum]['dent']:
@@ -109,7 +114,7 @@ def find_dir(inodes, inum, names, idx):
     return None
 
 
-def print_dent(ubifs, inodes, dent_node, long=True, longts=False):
+def print_dent(ubifs: Ubifs, inodes: Mapping[int, Inode], dent_node: nodes.dent_node, long: bool = True, longts: bool = False) -> None:
     inode = inodes[dent_node.inum]
     if long:
         fl = file_leng(ubifs, inode)
@@ -128,7 +133,7 @@ def print_dent(ubifs, inodes, dent_node, long=True, longts=False):
         print(dent_node.name)
 
 
-def file_leng(ubifs, inode):
+def file_leng(ubifs: Ubifs, inode: Inode) -> int:
     fl = 0
     if 'data' in inode:
         compr_type = 0

ubireader/ubifs/misc.py

@@ -17,6 +17,8 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #############################################################
 
+from __future__ import annotations
+from typing import TYPE_CHECKING, TypedDict
 from lzallright import LZOCompressor
 import struct
 import zlib
@@ -25,13 +27,22 @@
 from ubireader.debug import error
 from ubireader.ubifs.decrypt import lookup_inode_nonce, derive_key_from_nonce, datablock_decrypt
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from ubireader.ubifs import ubifs as Ubifs
+    from ubireader.ubifs.walk import Inode
+
 # For happy printing
 ino_types = ['file', 'dir','lnk','blk','chr','fifo','sock']
 node_types = ['ino','data','dent','xent','trun','pad','sb','mst','ref','idx','cs','orph']
 key_types = ['ino','data','dent','xent']
 
+class ParsedKey(TypedDict):
+    type: str
+    ino_num: int
+    khash: int
 
-def parse_key(key):
+def parse_key(key: bytes) -> ParsedKey:
     """Parse node key
 
     Arguments:
@@ -51,7 +62,7 @@ def parse_key(key):
     return {'type':key_type, 'ino_num':ino_num, 'khash': khash}
 
 
-def decompress(ctype, unc_len, data):
+def decompress(ctype: int, unc_len: int, data: bytes) -> bytes | None:
     """Decompress data.
 
     Arguments:
@@ -76,7 +87,7 @@ def decompress(ctype, unc_len, data):
         return data
 
 
-def process_reg_file(ubifs, inode, path, inodes):
+def process_reg_file(ubifs: Ubifs, inode: Inode, path: str, inodes: Mapping[int, Inode]) -> bytes:
     try:
         buf = bytearray()
         start_key = (UBIFS_DATA_KEY << UBIFS_S_KEY_BLOCK_BITS)