Implement optional caching of remote files

Author: AT0myks

README.md

@@ -90,7 +90,7 @@ the file name.
 #### Extract
 
 ```
-usage: reolinkfw extract [-h] [-d DEST] [-f] file_or_url
+usage: reolinkfw extract [-h] [--no-cache] [-d DEST] [-f] file_or_url
 
 Extract the file system from a Reolink firmware
 
@@ -99,6 +99,7 @@ positional arguments:
 
 optional arguments:
   -h, --help            show this help message and exit
+  --no-cache            don't use cache for remote files (URLs)
   -d DEST, --dest DEST  destination directory. Default: current directory
   -f, --force           overwrite existing files. Does not apply to UBIFS. Default: False
 ```

reolinkfw/__init__.py

@@ -21,7 +21,14 @@
 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, sha256_pak
+from reolinkfw.util import (
+    DummyLEB,
+    get_cache_file,
+    get_fs_from_ubi,
+    has_cache,
+    make_cache_file,
+    sha256_pak
+)
 
 __version__ = "1.1.0"
 
@@ -170,7 +177,7 @@ async def direct_download_url(url):
     return url
 
 
-async def get_paks(file_or_url) -> list[tuple[Optional[str], PAK]]:
+async def get_paks(file_or_url, 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
@@ -180,12 +187,16 @@ async def get_paks(file_or_url) -> list[tuple[Optional[str], PAK]]:
     It is the caller's responsibility to close the PAK files.
     """
     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))
         file_or_url = await direct_download_url(file_or_url)
         zip_or_pak_bytes = await download(file_or_url)
         if isinstance(zip_or_pak_bytes, int):
             raise Exception(f"HTTP error {zip_or_pak_bytes}")
-        elif is_pak_file(zip_or_pak_bytes):
-            pakname = dict(parse_qsl(urlparse(file_or_url).query)).get("name")
+        pakname = dict(parse_qsl(urlparse(file_or_url).query)).get("name")
+        if use_cache:
+            make_cache_file(file_or_url, zip_or_pak_bytes, pakname)
+        if is_pak_file(zip_or_pak_bytes):
             return [(pakname, PAK.from_bytes(zip_or_pak_bytes))]
         else:
             zipfile = io.BytesIO(zip_or_pak_bytes)
@@ -203,13 +214,13 @@ async def get_paks(file_or_url) -> list[tuple[Optional[str], PAK]]:
     raise Exception("Not a URL or file")
 
 
-async def get_info(file_or_url):
+async def get_info(file_or_url, use_cache: bool = True):
     """Retrieve firmware info from an on-disk file or a URL.
 
     The file or resource may be a ZIP or a PAK.
     """
     try:
-        paks = await get_paks(file_or_url)
+        paks = await get_paks(file_or_url, use_cache)
     except Exception as e:
         return [{"file": file_or_url, "error": str(e)}]
     if not paks:

reolinkfw/__main__.py

@@ -11,13 +11,13 @@
 from reolinkfw.util import sha256_pak
 
 
-def info(args: argparse.Namespace) -> None:
-    info = asyncio.run(get_info(args.file_or_url))
+async def info(args: argparse.Namespace) -> None:
+    info = await get_info(args.file_or_url, not args.no_cache)
     print(json.dumps(info, indent=args.indent, default=str))
 
 
 async def extract(args: argparse.Namespace) -> None:
-    paks = await get_paks(args.file_or_url)
+    paks = await get_paks(args.file_or_url, not args.no_cache)
     if not paks:
         raise Exception("No PAKs found in ZIP file")
     dest = Path.cwd() if args.dest is None else args.dest
@@ -32,24 +32,24 @@ def main():
     parser.add_argument("-V", "--version", action="version", version=f"%(prog)s {__version__}")
     subparsers = parser.add_subparsers(required=True)
 
-    parser_i = subparsers.add_parser("info")
+    pcache = argparse.ArgumentParser(add_help=False)
+    pcache.add_argument("--no-cache", action="store_true", help="don't use cache for remote files (URLs)")
+
+    parser_i = subparsers.add_parser("info", parents=[pcache])
     parser_i.add_argument("file_or_url", help="URL or on-disk file")
     parser_i.add_argument("-i", "--indent", type=int, help="indent level for pretty print")
     parser_i.set_defaults(func=info)
 
     descex = "Extract the file system from a Reolink firmware"
-    parser_e = subparsers.add_parser("extract", help=descex.lower(), description=descex)
+    parser_e = subparsers.add_parser("extract", parents=[pcache], help=descex.lower(), description=descex)
     parser_e.add_argument("file_or_url", help="URL or on-disk file")
     parser_e.add_argument("-d", "--dest", type=Path, help="destination directory. Default: current directory")
     parser_e.add_argument("-f", "--force", action="store_true", help="overwrite existing files. Does not apply to UBIFS. Default: %(default)s")
     parser_e.set_defaults(func=extract)
 
     args = parser.parse_args()
     try:
-        if asyncio.iscoroutinefunction(args.func):
-            asyncio.run(args.func(args))
-        else:
-            args.func(args)
+        asyncio.run(args.func(args))
     except Exception as e:
         sys.exit(f"error: {e}")
 

reolinkfw/util.py

@@ -2,14 +2,22 @@
 import io
 from contextlib import contextmanager
 from functools import partial
+from os import scandir
+from pathlib import Path
+from shutil import disk_usage
+from tempfile import gettempdir as _gettempdir
+from zipfile import is_zipfile
 
-from pakler import PAK
+from pakler import PAK, is_pak_file
 from ubireader.ubi import ubi
 from ubireader.ubi_io import ubi_file, leb_virtual_file
 from ubireader.utils import guess_peb_size
 
 from reolinkfw.tmpfile import TempFile
 
+ONEMIB = 1024**2
+ONEGIB = 1024**3
+
 
 class DummyLEB:
     """A class that emulates ubireader's `leb_virtual_file`."""
@@ -69,6 +77,57 @@ def get_fs_from_ubi(fd, size, offset=0) -> bytes:
 def sha256_pak(pak: PAK) -> str:
     sha = hashlib.sha256()
     pak._fd.seek(0)
-    for block in iter(partial(pak._fd.read, 1024**2), b''):
+    for block in iter(partial(pak._fd.read, ONEMIB), b''):
         sha.update(block)
     return sha.hexdigest()
+
+
+def dir_size(path):
+    size = 0
+    try:
+        with scandir(path) as it:
+            for entry in it:
+                if entry.is_dir(follow_symlinks=False):
+                    size += dir_size(entry.path)
+                elif entry.is_file(follow_symlinks=False):
+                    size += entry.stat().st_size
+    except OSError:
+        pass
+    return size
+
+
+def gettempdir() -> Path:
+    return Path(_gettempdir()) / "reolinkfwcache"
+
+
+def get_cache_file(url: str) -> Path:
+    file = gettempdir() / hashlib.sha256(url.encode("utf8")).hexdigest()
+    if is_zipfile(file) or is_pak_file(file):
+        return file
+    try:
+        with open(file, 'r', encoding="utf8") as f:
+            return gettempdir() / f.read(256)
+    except (OSError, UnicodeDecodeError):
+        return file
+
+
+def has_cache(url: str) -> bool:
+    return get_cache_file(url).is_file()
+
+
+def make_cache_file(url: str, filebytes, name=None) -> bool:
+    tempdir = gettempdir()
+    tempdir.mkdir(exist_ok=True)
+    if disk_usage(tempdir).free < ONEGIB or dir_size(tempdir) > ONEGIB:
+        return False
+    sha = hashlib.sha256(url.encode("utf8")).hexdigest()
+    name = sha if not isinstance(name, str) else name
+    try:
+        with open(tempdir / name, "wb") as f:
+            f.write(filebytes)
+        if name != sha:
+            with open(tempdir / sha, 'w', encoding="utf8") as f:
+                f.write(name)
+    except OSError:
+        return False
+    return True