More type hints

Author: AT0myks

reolinkfw/__init__.py

@@ -2,19 +2,22 @@
 import io
 import posixpath
 import re
+from collections.abc import Iterable, Mapping
 from pathlib import Path
-from typing import Optional
+from typing import IO, Any, BinaryIO, Optional, Union
 from urllib.parse import parse_qsl, urlparse
 from zipfile import ZipFile, is_zipfile
 
 import aiohttp
+from aiohttp.typedefs import StrOrURL
 from lxml.etree import fromstring
 from lxml.html import document_fromstring
 from pakler import PAK, Section, is_pak_file
 from pycramfs import Cramfs
 from PySquashfsImage import SquashFsImage
 
 from reolinkfw.tmpfile import TempFile
+from reolinkfw.typedefs import Buffer, Files, StrPath, StrPathURL
 from reolinkfw.ubifs import UBIFS
 from reolinkfw.uboot import get_arch_name, get_uboot_version, get_uimage_header
 from reolinkfw.util import (
@@ -35,7 +38,7 @@
 FS_SECTIONS = ROOTFS_SECTIONS + ["app"]
 
 
-async def download(url):
+async def download(url: StrOrURL) -> Union[bytes, int]:
     """Return resource as bytes.
 
     Return the status code of the request if it is not 200.
@@ -45,7 +48,7 @@ async def download(url):
             return await resp.read() if resp.status == 200 else resp.status
 
 
-def extract_paks(zip) -> list[tuple[str, PAK]]:
+def extract_paks(zip: Union[StrPath, IO[bytes]]) -> list[tuple[str, PAK]]:
     """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.
@@ -61,8 +64,8 @@ def extract_paks(zip) -> list[tuple[str, PAK]]:
     return paks
 
 
-def get_info_from_files(files):
-    xml = dict(fromstring(files["dvr.xml"]).items())
+def get_info_from_files(files: Mapping[Files, Optional[bytes]]) -> dict[str, Optional[str]]:
+    xml: dict[str, str] = dict(fromstring(files["dvr.xml"]).items())
     info = {k: xml.get(k) for k in INFO_KEYS}
     info["version_file"] = files["version_file"].decode().strip()
     if not info.get("firmware_version_prefix"):
@@ -72,7 +75,7 @@ def get_info_from_files(files):
     return info
 
 
-def get_files_from_squashfs(fd, offset=0, closefd=True):
+def get_files_from_squashfs(fd: BinaryIO, offset: int = 0, closefd: bool = True) -> dict[Files, Optional[bytes]]:
     # 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.
@@ -85,7 +88,7 @@ def get_files_from_squashfs(fd, offset=0, closefd=True):
     return files
 
 
-def get_files_from_ubifs(binbytes):
+def get_files_from_ubifs(binbytes: Buffer) -> dict[Files, Optional[bytes]]:
     # 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.
@@ -98,7 +101,7 @@ def get_files_from_ubifs(binbytes):
     return files
 
 
-def get_files_from_ubi(fd, size, offset=0):
+def get_files_from_ubi(fd: BinaryIO, size: int, offset: int = 0) -> dict[Files, Optional[bytes]]:
     fsbytes = get_fs_from_ubi(fd, size, offset)
     fs = FileType.from_magic(fsbytes[:4])
     if fs == FileType.UBIFS:
@@ -108,7 +111,7 @@ def get_files_from_ubi(fd, size, offset=0):
     raise Exception("Unknown file system in UBI")
 
 
-def get_files_from_cramfs(fd, offset=0, closefd=True):
+def get_files_from_cramfs(fd: BinaryIO, offset: int = 0, closefd: bool = True) -> dict[Files, Optional[bytes]]:
     # 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.
@@ -120,15 +123,15 @@ def get_files_from_cramfs(fd, offset=0, closefd=True):
     return files
 
 
-def is_url(string):
+def is_url(string: StrOrURL) -> bool:
     return str(string).startswith("http")
 
 
-def is_local_file(string):
+def is_local_file(string: StrPath) -> bool:
     return Path(string).is_file()
 
 
-def get_fs_info(pak: PAK, fs_sections: list[Section]) -> list[dict[str, str]]:
+def get_fs_info(pak: PAK, fs_sections: Iterable[Section]) -> list[dict[str, str]]:
     result = []
     for section in fs_sections:
         pak._fd.seek(section.start)
@@ -143,7 +146,7 @@ def get_fs_info(pak: PAK, fs_sections: list[Section]) -> list[dict[str, str]]:
     return result
 
 
-async def get_info_from_pak(pak: PAK):
+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]
     app = fs_sections[-1]
@@ -169,7 +172,7 @@ async def get_info_from_pak(pak: PAK):
     }
 
 
-async def direct_download_url(url):
+async def direct_download_url(url: str) -> str:
     if url.startswith("https://drive.google.com/file/d/"):
         return f"https://drive.google.com/uc?id={url.split('/')[5]}&confirm=t"
     elif url.startswith("https://www.mediafire.com/file/"):
@@ -182,7 +185,7 @@ async def direct_download_url(url):
     return url
 
 
-async def get_paks(file_or_url, 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], PAK]]:
     """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
@@ -191,6 +194,7 @@ async def get_paks(file_or_url, use_cache: bool = True) -> list[tuple[Optional[s
     be None. If the file is a ZIP the list might be empty.
     It is the caller's responsibility to close the PAK files.
     """
+    file_or_url = str(file_or_url)
     if is_url(file_or_url):
         if use_cache and has_cache(file_or_url):
             return await get_paks(get_cache_file(file_or_url))
@@ -219,7 +223,7 @@ async def get_paks(file_or_url, use_cache: bool = True) -> list[tuple[Optional[s
     raise Exception("Not a URL or file")
 
 
-async def get_info(file_or_url, use_cache: bool = True):
+async def get_info(file_or_url: StrPathURL, use_cache: bool = True) -> list[dict[str, Any]]:
     """Retrieve firmware info from an on-disk file or a URL.
 
     The file or resource may be a ZIP or a PAK.

reolinkfw/__main__.py

@@ -55,7 +55,7 @@ async def extract(args: Namespace) -> None:
         pakfile.close()
 
 
-def main():
+def main() -> None:
     parser = ArgumentParser(description="Extract information and files from Reolink firmwares")
     parser.add_argument("-V", "--version", action="version", version=f"%(prog)s {__version__}")
     subparsers = parser.add_subparsers(required=True)

reolinkfw/extract.py

@@ -1,6 +1,7 @@
 from contextlib import redirect_stdout
 from io import StringIO
 from pathlib import Path
+from typing import Optional
 
 from pakler import PAK, Section
 from pycramfs import Cramfs
@@ -17,7 +18,7 @@
 from reolinkfw.util import FileType, closing_ubifile, get_fs_from_ubi
 
 
-def extract_file_system(pak: PAK, section: Section, dest: Path = None):
+def extract_file_system(pak: PAK, 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)
@@ -47,7 +48,7 @@ def extract_file_system(pak: PAK, section: Section, dest: Path = None):
         raise Exception("Unknown file system")
 
 
-def extract_pak(pak: PAK, dest: Path = None, force: bool = False):
+def extract_pak(pak: PAK, 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]

reolinkfw/tmpfile.py

@@ -1,58 +1,62 @@
+from __future__ import annotations
+
 import os
 import platform
 import sys
 import tempfile
 from abc import abstractmethod
 
+from reolinkfw.typedefs import Buffer
+
 
 class TempFileFromBytes(os.PathLike):
 
-    def __init__(self, filebytes):
+    def __init__(self, filebytes: Buffer) -> None:
         self._filebytes = filebytes
         self._fd = -1
         self._path = None
 
-    def __enter__(self):
+    def __enter__(self) -> TempFileFromBytes:
         self.open()
         return self
-    
-    def __exit__(self, exc_type, exc_value, traceback):
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
         self.close()
-    
-    def __fspath__(self):
+
+    def __fspath__(self) -> str:
         return self._path
-    
+
     @abstractmethod
-    def open(self):
+    def open(self) -> int:
         ...
-    
+
     @abstractmethod
-    def close(self):
+    def close(self) -> None:
         ...
 
 
 class LinuxInMemoryFile(TempFileFromBytes):
 
-    def open(self):
+    def open(self) -> int:
         self._fd = os.memfd_create("temp")
         os.write(self._fd, self._filebytes)
         self._path = f"/proc/self/fd/{self._fd}"
         return self._fd
-    
-    def close(self):
+
+    def close(self) -> None:
         os.close(self._fd)
         self._fd = -1
         self._path = None
 
 
 class OnDiskTempFile(TempFileFromBytes):
-    
-    def open(self):
+
+    def open(self) -> int:
         self._fd, self._path = tempfile.mkstemp()
         os.write(self._fd, self._filebytes)
         return self._fd
 
-    def close(self):
+    def close(self) -> None:
         os.close(self._fd)
         os.unlink(self._path)
         self._fd = -1

reolinkfw/typedefs.py

@@ -0,0 +1,17 @@
+import collections.abc
+from os import PathLike
+from typing import AnyStr, Literal, Union
+
+from aiohttp.typedefs import StrOrURL
+
+if hasattr(collections.abc, "Buffer"):
+    Buffer = collections.abc.Buffer
+else:
+    Buffer = Union[bytes, bytearray, memoryview]
+
+Files = Literal["version_file", "version.json", "dvr.xml", "dvr", "router"]
+GenericPath = Union[AnyStr, PathLike[AnyStr]]
+StrPath = Union[str, PathLike[str]]
+StrPathURL = Union[StrPath, StrOrURL]
+StrOrBytesPath = Union[StrPath, bytes, PathLike[bytes]]
+FileDescriptorOrPath = Union[int, StrOrBytesPath]

reolinkfw/ubifs.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 from pathlib import PurePosixPath
 from stat import filemode
-from typing import Iterator, Literal, Optional
+from typing import Any, Iterator, Literal, Optional
 
 from ubireader.ubifs import ubifs, walk
 from ubireader.ubifs.defines import (
@@ -25,11 +26,12 @@
 from ubireader.utils import guess_leb_size
 
 from reolinkfw.tmpfile import TempFile
+from reolinkfw.typedefs import Buffer, FileDescriptorOrPath, StrPath
 
 
 class File:
 
-    def __init__(self, image: UBIFS, nodes: dict, name: str = '', parent: Optional[Directory] = None) -> None:
+    def __init__(self, image: UBIFS, nodes: Mapping[str, Any], name: str = '', parent: Optional[Directory] = None) -> None:
         self._image = image
         self._nodes = nodes
         self._name = name
@@ -108,9 +110,9 @@ def readlink(self) -> PurePosixPath:
 
 class Directory(File):
 
-    def __init__(self, image: UBIFS, nodes: dict, name: str = '', parent: Optional[Directory] = None) -> None:
+    def __init__(self, image: UBIFS, nodes: Mapping[str, Any], name: str = '', parent: Optional[Directory] = None) -> None:
         super().__init__(image, nodes, name, parent)
-        self._children = {}
+        self._children: dict[str, File] = {}
         for dent in nodes.get("dent", []):
             cls = filetype[dent.type]
             self._children[dent.name] = cls(image, image.inodes[dent.inum], dent.name, self)
@@ -126,7 +128,7 @@ def is_dir(self) -> Literal[True]:
     def children(self) -> dict[str, File]:
         return self._children
 
-    def select(self, path) -> Optional[File]:
+    def select(self, path: StrPath) -> Optional[File]:
         """Select a file of any kind by path.
 
         The path can be absolute or relative.
@@ -163,7 +165,7 @@ class UBIFS:
 
     def __init__(self, ubifs: ubifs) -> None:
         self._ubifs = ubifs
-        self._inodes = {}
+        self._inodes: dict[int, dict[str, Any]] = {}
         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])
@@ -178,7 +180,7 @@ def __iter__(self) -> Iterator[File]:
         yield from self._root.riter()
 
     @property
-    def inodes(self) -> dict:
+    def inodes(self) -> dict[int, dict[str, Any]]:
         return self._inodes
 
     @property
@@ -188,15 +190,15 @@ def root(self) -> Directory:
     def close(self) -> None:
         self._ubifs._file._fhandle.close()
 
-    def select(self, path) -> Optional[File]:
+    def select(self, path: StrPath) -> Optional[File]:
         return self._root.select(path)
 
     @classmethod
-    def from_file(cls, path) -> UBIFS:
+    def from_file(cls, path: FileDescriptorOrPath) -> UBIFS:
         return cls(ubifs(ubi_file(path, guess_leb_size(path))))
 
     @classmethod
-    def from_bytes(cls, bytes_, offset: int = 0) -> UBIFS:
+    def from_bytes(cls, bytes_: Buffer, 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")