From 60a885cb63f3b4bd27358230f198d1a7a8ec6206 Mon Sep 17 00:00:00 2001 From: m1stadev Date: Tue, 22 Aug 2023 14:30:16 -0500 Subject: [PATCH] [parser] Raise a RuntimeError if compression libs are not installed --- pyimg4/_parser.py | 32 ++++++++++++++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/pyimg4/_parser.py b/pyimg4/_parser.py index b7612f3..eb2f7b3 100644 --- a/pyimg4/_parser.py +++ b/pyimg4/_parser.py @@ -2,13 +2,25 @@ from zlib import adler32 import asn1 -import liblzfse -import lzss from Crypto.Cipher import AES from ._types import * from .errors import * +try: + import lzss + + have_lzss = True +except ImportError: + have_lzss = False + +try: + import liblzfse + + have_lzfse = True +except ImportError: + have_lzfse = False + class _PyIMG4: def __init__(self, data: Optional[bytes] = None) -> None: @@ -1247,12 +1259,20 @@ def compress(self, compression: Compression) -> None: ) if compression == Compression.LZSS: + if not have_lzss: + raise RuntimeError('pylzss not installed, cannot use LZSS compression') + self._data = self._create_complzss_header() + lzss.compress(self._data) if self.extra is not None: self._data += self.extra elif compression == Compression.LZFSE: + if not have_lzfse: + raise RuntimeError( + 'pyliblzfse not installed, cannot use LZFSE compression' + ) + payload_size = len(self._data) self._data = liblzfse.compress(self._data) # Cannot set LZFSE payload size until after compression @@ -1275,10 +1295,18 @@ def decompress(self) -> None: raise CompressionError('Cannot decompress encrypted payload.') elif self.compression == Compression.LZSS: + if not have_lzss: + raise RuntimeError('pylzss not installed, cannot use LZSS compression') + self._parse_complzss_header() self._data = lzss.decompress(self._data) elif self.compression == Compression.LZFSE: + if not have_lzfse: + raise RuntimeError( + 'pyliblzfse not installed, cannot use LZFSE compression' + ) + self._lzfse_payload_size = None self._data = liblzfse.decompress(self._data)