Replace is_* functions by an enum

Author: AT0myks

reolinkfw/__init__.py

@@ -2,6 +2,7 @@
 import io
 import posixpath
 import re
+from enum import Enum
 from pathlib import Path
 from typing import Optional
 from urllib.parse import parse_qsl, urlparse
@@ -12,19 +13,15 @@
 from lxml.html import document_fromstring
 from pakler import PAK, is_pak_file
 from pycramfs import Cramfs
+from pycramfs.const import MAGIC_BYTES as CRAMFS_MAGIC
 from PySquashfsImage import SquashFsImage
+from PySquashfsImage.const import SQUASHFS_MAGIC
+from ubireader.ubi.defines import UBI_EC_HDR_MAGIC as UBI_MAGIC
 from ubireader.ubifs import ubifs, walk
+from ubireader.ubifs.defines import UBIFS_NODE_MAGIC as UBIFS_MAGIC
 from ubireader.ubifs.output import _process_reg_file
 
-from reolinkfw.util import (
-    DummyLEB,
-    get_fs_from_ubi,
-    is_cramfs,
-    is_squashfs,
-    is_ubi,
-    is_ubifs,
-    sha256_pak
-)
+from reolinkfw.util import DummyLEB, get_fs_from_ubi, sha256_pak
 
 __version__ = "1.1.0"
 
@@ -35,6 +32,20 @@
 FS_SECTIONS = ROOTFS_SECTIONS + ["app"]
 
 
+class FileSystem(Enum):
+    CRAMFS = CRAMFS_MAGIC
+    SQUASHFS = SQUASHFS_MAGIC.to_bytes(4, "little")
+    UBI = UBI_MAGIC  # Not a file system
+    UBIFS = UBIFS_MAGIC
+
+    @classmethod
+    def from_magic(cls, key, default=None):
+        try:
+            return cls(key)
+        except ValueError:
+            return default
+
+
 async def download(url):
     """Return resource as bytes.
 
@@ -111,12 +122,12 @@ def get_files_from_ubifs(binbytes):
 
 def get_files_from_ubi(binbytes):
     fsbytes = get_fs_from_ubi(binbytes)
-    if is_ubifs(fsbytes):
+    fs = FileSystem.from_magic(fsbytes[:4])
+    if fs == FileSystem.UBIFS:
         return get_files_from_ubifs(fsbytes)
-    elif is_squashfs(fsbytes):
+    elif fs == FileSystem.SQUASHFS:
         return get_files_from_squashfs(fsbytes)
-    else:
-        raise Exception("unknown file system in UBI")
+    raise Exception("Unknown file system in UBI")
 
 
 def get_files_from_cramfs(binbytes):
@@ -144,13 +155,12 @@ async def get_info_from_pak(pak: PAK):
     binbytes = await asyncio.to_thread(extract_fs, pak)
     if isinstance(binbytes, str):
         return {"error": binbytes, "sha256": ha}
-    if is_cramfs(binbytes):
-        func = get_files_from_cramfs
-    elif is_ubi(binbytes):
-        func = get_files_from_ubi
-    elif is_squashfs(binbytes):
-        func = get_files_from_squashfs
-    else:
+    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:
         return {"error": "Unrecognized image type", "sha256": ha}
     files = await asyncio.to_thread(func, binbytes)
     info = get_info_from_files(files)

reolinkfw/extract.py

@@ -10,31 +10,25 @@
 from ubireader.ubifs import ubifs
 from ubireader.ubifs.output import extract_files as extract_ubifs
 
-from reolinkfw import FS_SECTIONS, ROOTFS_SECTIONS
-from reolinkfw.util import (
-    DummyLEB,
-    get_fs_from_ubi,
-    is_cramfs,
-    is_squashfs,
-    is_ubi,
-    is_ubifs
-)
+from reolinkfw import FS_SECTIONS, ROOTFS_SECTIONS, FileSystem
+from reolinkfw.util import DummyLEB, get_fs_from_ubi
 
 
 def extract_file_system(fs_bytes, dest: Path = None):
     dest = (Path.cwd() / "reolink_fs") if dest is None else dest
     dest.mkdir(parents=True, exist_ok=True)
-    if is_ubi(fs_bytes):
+    fs = FileSystem.from_magic(fs_bytes[:4])
+    if fs == FileSystem.UBI:
         extract_file_system(get_fs_from_ubi(fs_bytes), dest)
-    elif is_ubifs(fs_bytes):
+    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)
-    elif is_squashfs(fs_bytes):
+    elif fs == FileSystem.SQUASHFS:
         with SquashFsImage.from_bytes(fs_bytes) as image:
             extract_squashfs(image.root, dest, True)
-    elif is_cramfs(fs_bytes):
+    elif fs == FileSystem.CRAMFS:
         with Cramfs.from_bytes(fs_bytes) as image:
             extract_cramfs(image.rootdir, dest, True)
     else:

reolinkfw/util.py

@@ -4,18 +4,12 @@
 from functools import partial
 
 from pakler import PAK
-from pycramfs.const import MAGIC_BYTES as CRAMFS_MAGIC
-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.ubifs.defines import UBIFS_NODE_MAGIC as UBIFS_MAGIC
 from ubireader.utils import guess_peb_size
 
 from reolinkfw.tmpfile import TempFile
 
-SQUASHFS_MAGIC = SQUASHFS_MAGIC.to_bytes(4, "little")
-
 
 class DummyLEB:
     """A class that emulates ubireader's `leb_virtual_file`."""
@@ -70,22 +64,6 @@ def get_fs_from_ubi(binbytes):
             return b''.join(leb_virtual_file(ubi_obj, vol_blocks).reader())
 
 
-def is_ubi(bytes_):
-    return bytes_[:4] == UBI_MAGIC
-
-
-def is_squashfs(bytes_):
-    return bytes_[:4] == SQUASHFS_MAGIC
-
-
-def is_cramfs(bytes_):
-    return bytes_[:4] == CRAMFS_MAGIC
-
-
-def is_ubifs(bytes_):
-    return bytes_[:4] == UBIFS_MAGIC
-
-
 def sha256_pak(pak: PAK) -> str:
     sha = hashlib.sha256()
     pak._fd.seek(0)