add encrption module

fix bug in zstd compression warp code
fix bug when showing compression inforamtion

game_backuper/_zstd.pyi

@@ -0,0 +1,15 @@
+def version() -> str: ...
+def maxCLevel() -> int: ...
+class ZSTDCompressor:  # noqa: E302
+    def __init__(self, compresslevel: int = 3): ...
+    def compress(self, inp: bytes) -> bytes: ...
+    def flush(self) -> bytes: ...
+class ZSTDDecompressor:  # noqa: E302
+    def __init__(self): ...
+    def decompress(self, data: bytes, max_length: int = -1) -> bytes: ...
+    @property
+    def eof(self) -> bool: ...
+    @property
+    def unused_data(self) -> bytes: ...
+    @property
+    def needs_input(self) -> bool: ...

game_backuper/_zstd.pyx

@@ -174,7 +174,7 @@ cdef class ZSTDDecompressor:
         if self.finish and i.pos < i.size:
             obuf = (<char*> i.src) + i.pos
             self._unused_data = PyBytes_FromStringAndSize(obuf, i.size - i.pos)
-        if max_length == 1 or len(b) <= max_length:
+        if max_length == -1 or len(b) <= max_length:
             return b
         else:
             self._buff = b[max_length:]

game_backuper/compress.py

@@ -1,5 +1,5 @@
 try:
-    from bz2 import BZ2File
+    from bz2 import BZ2File, BZ2Compressor, BZ2Decompressor
     have_bz2 = True
 except ImportError:
     have_bz2 = False
@@ -9,7 +9,7 @@ try:
 except ImportError:
     have_gzip = False
 try:
-    from lzma import LZMAFile
+    from lzma import LZMAFile, LZMACompressor, LZMADecompressor
     have_lzma = True
 except ImportError:
     have_lzma = False
@@ -17,14 +17,21 @@ try:
     from lzip import (
         FileEncoder as LZIPFileEncoder,
         decompress_file_iter as LZIP_decompress_file_iter,
+        level_to_dictionary_size_and_match_len_limit as lzip_level_dict,
+    )
+    from lzip_extension import (
+        Decoder as LZIPDecoder,
+        Encoder as LZIPEncoder,
     )
     have_lzip = True
 except ImportError:
     have_lzip = False
 try:
     from game_backuper.zstd import (
+        ZSTDCompressor,
+        ZSTDDecompressor,
         ZSTDFile,
-        MAX_COMPRESS_LEVEL as ZSTD_MAX
+        MAX_COMPRESS_LEVEL as ZSTD_MAX,
     )
     have_zstd = True
 except ImportError:
@@ -55,6 +62,87 @@ from os import remove
 from game_backuper.file import mkdir_for_file, hydrate_file_if_needed
 
 
+if have_gzip:
+    class GZCompressor:
+        def __init__(self, compress_level: int = 9, file_obj=None):
+            self.write_to_file = True
+            self._fileobj = file_obj
+            self._level = compress_level
+            self._f = None
+
+        def close(self):
+            self._f.close()
+
+        def compress(self, data: bytes) -> bytes:
+            if self._f is None:
+                self._f = GzipFile(mode="wb", compresslevel=self._level, fileobj=self._fileobj)  # noqa: E501
+            self._f.write(data)
+            return b''
+
+        def flush(self) -> bytes:
+            self._f.flush()
+            return b''
+
+    class GZDecompressor:
+        def __init__(self, file_obj=None):
+            self.write_to_file = True
+            self._fileobj = file_obj
+            self._f = None
+
+        def close(self):
+            self._f.close()
+
+        def read(self, len: int) -> bytes:
+            if self._f is None:
+                self._f = GzipFile(mode="rb", fileobj=self._fileobj)
+            return self._f.read(len)
+
+if have_lzip:
+    class LZIPCompressor:
+        def __init__(self, level: int = 6):
+            d = lzip_level_dict[level]
+            self._encoder = LZIPEncoder(d[0], d[1], 1 << 51)
+
+        def compress(self, data: bytes) -> bytes:
+            return self._encoder.compress(data)
+
+        def flush(self) -> bytes:
+            return self._encoder.finish()
+
+    class LZIPDecompressor:
+        def __init__(self):
+            self._decoder = LZIPDecoder(1)
+
+        def decompress(self, data: bytes) -> bytes:
+            return self._decoder.decompress(data)
+
+        @property
+        def eof(self):
+            return False
+
+if have_brotli:
+    class BrotliCompressor2:
+        def __init__(self, *k, **kw):
+            self._c = BrotliCompressor(*k, **kw)
+
+        def compress(self, data: bytes) -> bytes:
+            return self._c.process(data)
+
+        def flush(self) -> bytes:
+            return self._c.finish()
+
+    class BrotliDecompressor2:
+        def __init__(self, *k, **kw) -> None:
+            self._c = BrotliDecompressor(*k, **kw)
+
+        def decompress(self, data: bytes) -> bytes:
+            return self._c.process(data)
+
+        @property
+        def eof(self):
+            return self._c.is_finished()
+
+
 @unique
 class CompressMethod(IntEnum):
     BZIP2 = 0
@@ -169,6 +257,41 @@ class CompressConfig:
         t = type(self)
         return f"<{t.__module__}.{t.__qualname__} object at {hex(id(self))}; method={repr(self._method)}, level={repr(self._level)}, ext={repr(self._ext)}>"  # noqa: E501
 
+    def compressor(self, fileobj):
+        if self._method == CompressMethod.BZIP2:
+            return BZ2Compressor(self._level)
+        elif self._method == CompressMethod.GZIP:
+            return GZCompressor(self._level, fileobj)
+        elif self._method == CompressMethod.LZMA:
+            return LZMACompressor(preset=self._level)
+        elif self._method == CompressMethod.LZIP:
+            return LZIPCompressor(self._level)
+        elif self._method == CompressMethod.ZSTD:
+            return ZSTDCompressor(self._level)
+        elif self._method == CompressMethod.SNAPPY:
+            return Snappy_Compressor()
+        elif self._method == CompressMethod.BROTLI:
+            c = {}
+            if self._level is not None:
+                c['quality'] = self._level
+            return BrotliCompressor2(**c)
+
+    def decompressor(self, fileobj):
+        if self._method == CompressMethod.BZIP2:
+            return BZ2Decompressor()
+        elif self._method == CompressMethod.GZIP:
+            return GZDecompressor(fileobj)
+        elif self._method == CompressMethod.LZMA:
+            return LZMADecompressor()
+        elif self._method == CompressMethod.LZIP:
+            return LZIPDecompressor()
+        elif self._method == CompressMethod.ZSTD:
+            return ZSTDDecompressor()
+        elif self._method == CompressMethod.SNAPPY:
+            return Snappy_Decompressor()
+        elif self._method == CompressMethod.BROTLI:
+            return BrotliDecompressor2()
+
     @cached_property
     def chunk_size(self) -> int:
         return self._chunk_size
@@ -212,6 +335,8 @@ def sizeof_fmt(num, suffix='B'):
 
 
 def compress_info(ori: int, re: int):
+    if ori == 0:
+        return f"{sizeof_fmt(ori)} -> {sizeof_fmt(re)} ({float('inf')})"
     return f"{sizeof_fmt(ori)} -> {sizeof_fmt(re)} ({re/ori*100:.2f}%)"
 
 

game_backuper/enc.py

@@ -0,0 +1,309 @@
+from base64 import b85decode
+from io import RawIOBase
+from os import PathLike, urandom
+from typing import Union
+from zlib import crc32
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from game_backuper.compress import CompressConfig
+
+
+_MODE_CLOSED = 0
+_MODE_READ = 1
+_MODE_WRITE = 3
+
+
+class DecryptException(Exception):
+    pass
+
+
+class EncFile(RawIOBase):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, val, tb):
+        if not self.closed:
+            self.close()
+
+    def __init__(self, fn, mode: str, salt: Union[bytes, str],
+                 key: Union[bytes, str] = None, iv: Union[bytes, str] = None,
+                 length: int = None, crc32: Union[str, int] = None,
+                 compress: CompressConfig = None):
+        self._fp = None
+        self._mode = _MODE_CLOSED
+        if mode in ["", "r", "rb"]:
+            mode = "rb"
+            mode_code = _MODE_READ
+        elif mode in ["w", "wb"]:
+            mode = "wb"
+            mode_code = _MODE_WRITE
+        elif mode in ["x", "xb"]:
+            mode = "xb"
+            mode_code = _MODE_WRITE
+        elif mode in ["a", "ab"]:
+            mode = "ab"
+            mode_code = _MODE_WRITE
+        if isinstance(fn, (str, bytes, PathLike)):
+            self._fp = open(fn, mode)
+            self._mode = mode_code
+        else:
+            raise TypeError('filename must be str, bytes or PathLike.')
+        if self._mode == _MODE_WRITE:
+            self._key = urandom(32)
+            self._iv = urandom(16)
+            self._cipher = Cipher(algorithms.AES(self._key), modes.CBC(self._iv))  # noqa: E501
+            self._enc = self._cipher.encryptor()
+            self._crc32 = 0
+            if compress is not None:
+                self._compressor = compress.compressor(self)
+                if self._compressor is None:
+                    raise NotImplementedError('Unsupported compression type.')
+            else:
+                self._compressor = None
+        if isinstance(salt, str):
+            self._salt = b85decode(salt)
+        else:
+            self._salt = salt
+        if self._mode == _MODE_READ:
+            if isinstance(key, str):
+                key = b85decode(key)
+            if isinstance(iv, str):
+                iv = b85decode(iv)
+            if key is None or len(key) != 32:
+                raise ValueError('A 256-bit key is required.')
+            self._key = key
+            if iv is None or len(iv) != 16:
+                raise ValueError('A 128-bit initialization_vector is required.')  # noqa: E501
+            self._iv = iv
+            self._cipher = Cipher(algorithms.AES(self.key), modes.CBC(self._iv))  # noqa: E501
+            self._dec = self._cipher.decryptor()
+            if length is None or not isinstance(length, int):
+                raise ValueError("data's length is needed.")
+            self._length = length
+            self._buf = b''
+            self._eof = False
+            if crc32:
+                if isinstance(crc32, str):
+                    self._crc32 = int(crc32, 16)
+                else:
+                    self._crc32 = crc32
+                self._crc32_checked = False
+            else:
+                self._crc32 = None
+            self._decompressor = None
+            if compress:
+                self._decompressor = compress.decompressor(self)
+                if self._decompressor is None:
+                    raise NotImplementedError('Unsupported compression type.')
+                self._debuf = b''
+        self._pos = 0
+        self._crc_size = algorithms.AES.block_size * 8
+        self._flushing = False
+
+    def _check_not_closed(self):
+        if self.closed:
+            raise ValueError("I/O operation on closed file")
+
+    def close(self):
+        if self._mode == _MODE_CLOSED:
+            return
+        try:
+            if self._mode == _MODE_READ:
+                if hasattr(self._decompressor, 'close'):
+                    self._decompressor.write_to_file = False
+                    self._decompressor.close()
+                    self._decompressor.write_to_file = True
+            if self._mode == _MODE_WRITE:
+                if hasattr(self._compressor, 'close'):
+                    self._compressor.write_to_file = False
+                    self._compressor.close()
+                    self._compressor.write_to_file = True
+                elif self._compressor:
+                    d = self._compressor.flush()
+                    if d:
+                        length = len(d)
+                        if self._pos < self._crc_size:
+                            le = min(length, self._crc_size - self._pos)
+                            self._crc32 = crc32(d[:le], self._crc32)
+                        self._fp.write(self._enc.update(d))
+                        self._pos += length
+                if self._pos % algorithms.AES.block_size != 0:
+                    self._fp.write(self._enc.update(b"\x00" * (algorithms.AES.block_size - (self._pos % algorithms.AES.block_size))))  # noqa: E501
+                self._fp.write(self._enc.finalize())
+        finally:
+            self._fp.close()
+            self._dec = None
+            self._enc = None
+            self._cipher = None
+            self._fp = None
+            self._mode = _MODE_CLOSED
+
+    @property
+    def closed(self):
+        return self._mode == _MODE_CLOSED
+
+    def check_crc32(self):
+        while len(self._buf) < self._crc_size:
+            if not self.__read():
+                break
+        if crc32(self._buf[:min(self._crc_size, self._length)]) == self._crc32:
+            return True
+        else:
+            return False
+
+    @property
+    def crc32(self):
+        return '{:08x}'.format(self._crc32)
+
+    def __decompress(self):
+        if len(self._buf) == 0:
+            return False
+        if self._decompressor:
+            le = min(len(self._buf), self._length - self._pos)
+            if le == 0 or (hasattr(self._decompressor, "eof") and self._decompressor.eof):  # noqa: E501
+                return False
+            self._debuf += self._decompressor.decompress(self._buf[:le])
+            self._pos += le
+            self._buf = self._buf[le:]
+            return True
+        return False
+
+    @property
+    def eof(self):
+        return self._pos >= self._length
+
+    def fileno(self):
+        self._check_not_closed()
+        return self._fp.fileno()
+
+    def flush(self) -> None:
+        if self._flushing:
+            return
+        self._flushing = True
+        self._check_not_closed()
+        if hasattr(self._compressor, 'write_to_file'):
+            self._compressor.write_to_file = False
+            self._compressor.flush()
+            self._compressor.write_to_file = True
+        self._fp.flush()
+        self._flushing = False
+
+    def readable(self):
+        self._check_not_closed()
+        return self._mode == _MODE_READ
+
+    def __read(self):
+        data = self._fp.read(algorithms.AES.block_size)
+        if not data:
+            if not self._eof:
+                self._buf += self._dec.finalize()
+                self._eof = True
+            else:
+                return False
+        else:
+            self._buf += self._dec.update(data)
+        return True
+
+    def read(self, size: int = -1):
+        if not self.readable():
+            raise ValueError('File is not readable.')
+        if self._crc32 is not None and not self._crc32_checked:
+            if not self.check_crc32():
+                raise DecryptException("crc32 check failed.")
+            self._crc32_checked = True
+        if size < 0:
+            return self.readall()
+        if self._decompressor and hasattr(self._decompressor, 'write_to_file') and self._decompressor.write_to_file:  # noqa: E501
+            self._decompressor.write_to_file = False
+            d = self._decompressor.read(size)
+            self._decompressor.write_to_file = True
+            return d
+        if self._decompressor and not hasattr(self._decompressor, 'write_to_file'):  # noqa: E501
+            if not size or (self.eof and len(self._debuf) == 0):
+                return b""
+            while True:
+                if not self.__read():
+                    self.__decompress()
+                    break
+                if not self.__decompress():
+                    break
+                if size <= len(self._debuf):
+                    data = self._debuf[:size]
+                    self._debuf = self._debuf[size:]
+                    return data
+            d = self._debuf[:size]
+            self._debuf = self._debuf[size:]
+            return d
+        if not size or self.eof:
+            return b""
+        size = min(size, self._length - self._pos)
+        if size <= len(self._buf):
+            data = self._buf[:size]
+            self._buf = self._buf[size:]
+            self._pos += size
+            return data
+        while True:
+            if not self.__read():
+                break
+            if size <= len(self._buf):
+                data = self._buf[:size]
+                self._buf = self._buf[size:]
+                self._pos += size
+                return data
+        self._pos += min(len(self._buf), size)
+        d = self._buf[:size]
+        self._buf = self._buf[size:]
+        return d
+
+    def readinto(self, b):
+        with memoryview(b) as view, view.cast("B") as byte_view:
+            data = self.read(len(byte_view))
+            byte_view[:len(data)] = data
+        return len(data)
+
+    def tell(self):
+        return self._pos
+
+    def writable(self):
+        self._check_not_closed()
+        return self._mode == _MODE_WRITE
+
+    def write(self, data):
+        if not self.writable():
+            raise ValueError("File was not opened for writing")
+        if isinstance(data, bytes):
+            length = len(data)
+        elif isinstance(data, str):
+            data = data.encode()
+            length = len(data)
+        elif isinstance(data, bytearray):
+            data = bytes(data)
+            length = len(data)
+        else:
+            data = memoryview(data)
+            length = data.nbytes
+            data = data.tobytes()
+        if self._compressor:
+            if hasattr(self._compressor, 'write_to_file'):
+                if self._compressor.write_to_file:
+                    self._compressor.write_to_file = False
+                    self._compressor.compress(data)
+                    self._compressor.write_to_file = True
+                    return
+            else:
+                data = self._compressor.compress(data)
+                length = len(data)
+        if self._pos < self._crc_size:
+            le = min(length, self._crc_size - self._pos)
+            self._crc32 = crc32(data[:le], self._crc32)
+        self._fp.write(self._enc.update(data))
+        self._pos += length
+
+    @property
+    def key(self):
+        if len(self._salt) < 32:
+            self._salt = self._salt + b'\x00' * (32 - len(self._salt))
+        return bytes(a ^ b for a, b in zip(self._salt, self._key))
+
+    @property
+    def iv(self):
+        return self._iv

game_backuper/file.py

@@ -15,7 +15,11 @@ else:
     have_cfapi = False
 
 
-File = namedtuple('File', ['id', 'file', 'size', 'program', 'hash', 'type'])
+_File = namedtuple('File', ['id', 'file', 'size', 'program', 'hash', 'type'])
+
+
+class File(_File):
+    pass
 
 
 def mkdir_for_file(p: str):

testenc.py

@@ -0,0 +1,100 @@
+from game_backuper.compress import CompressConfig
+from game_backuper.enc import DecryptException, EncFile
+from os import urandom, remove
+from hashlib import sha512
+from zlib import crc32
+
+
+datalen = 4096
+data = urandom(datalen)
+a = sha512()
+a.update(data)
+with EncFile('a.txt', 'wb', a.digest()) as f:
+    f.write(data)
+    key = f.key
+    iv = f.iv
+    crc = f.crc32
+    print(crc)
+    crc2 = '{:08x}'.format(crc32(data[:1024]))
+    print(crc2)
+    assert crc == crc2
+with EncFile('a.txt', 'rb', a.digest(), key, iv, datalen, crc) as f:
+    d = f.read()
+    assert d == data
+with EncFile('a.txt', 'rb', b'', key, iv, datalen, crc) as f:
+    try:
+        d = f.read()
+        assert False
+    except DecryptException:
+        pass
+remove('a.txt')
+with EncFile('a.txt', 'wb', b'', compress=CompressConfig('gzip', 9)) as f:
+    f.write(data)
+    f.flush()
+    key = f.key
+    iv = f.iv
+le = f.tell()
+crc = f.crc32
+with EncFile('a.txt', 'rb', b'', key, iv, le, crc, compress=CompressConfig('gzip')) as f:  # noqa: E501
+    assert data == f.read()
+remove('a.txt')
+with EncFile('a.txt', 'wb', b'', compress=CompressConfig('bzip2', 1)) as f:
+    f.write(data)
+    f.flush()
+    key = f.key
+    iv = f.iv
+le = f.tell()
+crc = f.crc32
+with EncFile('a.txt', 'rb', b'', key, iv, le, crc, CompressConfig('bzip2', 1)) as f:  # noqa: E501
+    assert data == f.read()
+remove('a.txt')
+with EncFile('a.txt', 'wb', b'', compress=CompressConfig('lzma', 1)) as f:
+    f.write(data)
+    f.flush()
+    key = f.key
+    iv = f.iv
+le = f.tell()
+crc = f.crc32
+with EncFile('a.txt', 'rb', b'', key, iv, le, crc, CompressConfig('lzma', 1)) as f:  # noqa: E501
+    assert data == f.read()
+remove('a.txt')
+with EncFile('a.txt', 'wb', b'', compress=CompressConfig('lzip', 1)) as f:
+    f.write(data)
+    f.flush()
+    key = f.key
+    iv = f.iv
+le = f.tell()
+crc = f.crc32
+with EncFile('a.txt', 'rb', b'', key, iv, le, crc, CompressConfig('lzip', 1)) as f:  # noqa: E501
+    assert data == f.read()
+remove('a.txt')
+with EncFile('a.txt', 'wb', b'', compress=CompressConfig('zstd', 1)) as f:
+    f.write(data)
+    f.flush()
+    key = f.key
+    iv = f.iv
+le = f.tell()
+crc = f.crc32
+with EncFile('a.txt', 'rb', b'', key, iv, le, crc, CompressConfig('zstd', 1)) as f:  # noqa: E501
+    assert data == f.read()
+remove('a.txt')
+with EncFile('a.txt', 'wb', b'', compress=CompressConfig('snappy', 1)) as f:
+    f.write(data)
+    f.flush()
+    key = f.key
+    iv = f.iv
+le = f.tell()
+crc = f.crc32
+with EncFile('a.txt', 'rb', b'', key, iv, le, crc, CompressConfig('snappy', 1)) as f:  # noqa: E501
+    assert data == f.read()
+remove('a.txt')
+with EncFile('a.txt', 'wb', b'', compress=CompressConfig('brotli', 1)) as f:
+    f.write(data)
+    f.flush()
+    key = f.key
+    iv = f.iv
+le = f.tell()
+crc = f.crc32
+with EncFile('a.txt', 'rb', b'', key, iv, le, crc, CompressConfig('brotli', 1)) as f:  # noqa: E501
+    assert data == f.read()
+remove('a.txt')