def zstdlines(filename, encoding="utf-8", bufsize=65536): """ Generator over lines from a zstd compressed file. >>> for line in zstdlines("file.zst"): ... print(line) """ with open(filename, "rb") as f: decomp = ZstdDecompressor() with decomp.stream_reader(f) as reader: prev_line = "" while True: chunk = reader.read(bufsize) if not chunk: break while True: # We start with bytes but want unicode, which might not # align; so we jitter around the end to complete the # codepoint. try: string_data = chunk.decode(encoding) except UnicodeDecodeError: chunk = chunk + reader.read(1) else: break lines = string_data.split("\n") for i, line in enumerate(lines[:-1]): if i == 0: line = prev_line + line yield line prev_line = lines[-1]
def run(judge_class: Callable[[], JudgeDriver], task: JudgeTask) -> JudgeStatus: LOGGER.info('judge start (contest_id: {}, problem_id: {}, ' 'submission_id: {}, user_id: {}'.format( task.contest_id, task.problem_id, task.id, task.user_id)) zctx = ZstdDecompressor() try: task.code = zctx.decompress(task.code) for test in task.tests: test.input = zctx.decompress(test.input) test.output = zctx.decompress(test.output) except Exception: LOGGER.warning('decompress failed', exc_info=True) with transaction() as s: return _update_submission_status(s, task, JudgeStatus.InternalError) with judge_class() as judge: ret = _prepare(judge, task) if ret: return ret if task.compile_image_name: ret = _compile(judge, task) if ret: return ret ret = _tests(judge, task) LOGGER.info('judge finished (submission_id={}): {}'.format(task.id, ret)) return ret
def __enter__(self): codec = detect_compression(self.file) # (1) Open the input file # We keep a reference (`f`) to the original file, # in order to be able to close it on exit. # Calling `close()` on `ZstdDecompressor` does not # close the underlying resource. self.f = self.file.open("rb") # UTF-8 decoded, decompressed, file self.fb = self.f # (2) Setup the decompressor, if needed if codec == CompressionFormat.Zstandard: dict_data = ZstdCompressionDict(dictionary.read_bytes()) ctx = ZstdDecompressor(dict_data=dict_data) self.fb = ctx.stream_reader(self.fb, read_across_frames=True) # (3) Decode the file self.fb = TextIOWrapper(self.fb, "utf-8") # (4) Deserialize the records stream = map(json_tryloads, self.fb) # (5) Apply the filters stream = filter( lambda record: all(fn.keep(record) for fn in self.filters), stream) # (6) Apply the transformers for fn in self.transformers: stream = map(fn, stream) return stream
def test_log(tmpfile): records = [{ "msm_id": 1234, "prb_id": 5678 }, { "msm_id": 9876, "prb_id": 5432 }] with AtlasRecordsWriter(tmpfile, compression=True, log=True) as w: log_file = w.log_file for record in records: w.write(record) # TODO: Methods to simplify log reading? # Zstandard decompression context dict_data = ZstdCompressionDict(dictionary.read_bytes()) ctx = ZstdDecompressor(dict_data=dict_data) f = tmpfile.open("rb") log_f = log_file.open("rb") log = LogEntry.iter_unpack(log_f.read()) for i, (size, msm_id, prb_id) in enumerate(log): rec = json.loads(ctx.decompress(f.read(size)).decode("utf-8")) assert rec == records[i] assert msm_id == records[i]["msm_id"] assert prb_id == records[i]["prb_id"] f.close() log_f.close()
def __init__(self, item_type: type, index_path: str, num_pages: int, page_size: int): super().__init__(item_type, num_pages, page_size) self.index_path = index_path self.compressor = ZstdCompressor() self.decompressor = ZstdDecompressor() self.index_file = None self.mmap = None
def __decompress_zst_file(self, file, with_extension): with open(file, 'rb') as compressed: decomp = ZstdDecompressor() filename = os.path.splitext(file)[0] file_name = f"{self.destination_folder}/{file.name}{with_extension}" with open(file_name, 'wb') as destination: decomp.copy_stream(compressed, destination) print_green(f"unpacked zst file completed to {file_name}")
def zopen(fn, mode="r", *args, **kwargs): import codecs objs = None if fn.endswith(".gz"): import gzip objs = [gzip.open(fn, "rb")] elif fn.endswith(".bz2"): import bz2 objs = [bz2.open(fn, "rb")] elif fn.endswith(".xz"): import lzma objs = [lzma.open(fn, "rb")] elif fn.endswith(".zst"): from zstandard import ZstdDecompressor try: # documentation says KiB but it seems to be bytes ctx = ZstdDecompressor(max_window_size=1024 * 1024 * 1024 * 2) except: # fallback in case that changes ctx = ZstdDecompressor(max_window_size=1024 * 1024 * 2) f1 = open(fn, "rb", 512 * 1024) f2 = ctx.stream_reader(f1) objs = [f2, f1] else: objs = [open(fn, "rb", 512 * 1024)] if "b" not in mode: enc = kwargs.get("encoding", "utf-8") # yield io.TextIOWrapper(io.BufferedReader(f2)) yield codecs.getreader(enc)(objs[0]) else: yield objs[0] for obj in objs: obj.close()
def open_compressed(path): if is_gzip(path): return gzip.open(path) if is_zstd(path): from zstandard import ZstdDecompressor return ZstdDecompressor().stream_reader(open(path, 'rb')) return open(path, 'rb')
def __decompressNcz(nspf, f): ncaHeaderSize = 0x4000 blockID = 0 nspf.seek(0) header = nspf.read(ncaHeaderSize) start = f.tell() magic = nspf.read(8) if not magic == b'NCZSECTN': raise ValueError("No NCZSECTN found! Is this really a .ncz file?") sectionCount = readInt64(nspf) sections = [Section(nspf) for _ in range(sectionCount)] nca_size = ncaHeaderSize for i in range(sectionCount): nca_size += sections[i].size pos = nspf.tell() blockMagic = nspf.read(8) nspf.seek(pos) useBlockCompression = blockMagic == b'NCZBLOCK' blockSize = -1 if useBlockCompression: BlockHeader = Block(nspf) blockDecompressorReader = BlockDecompressorReader.BlockDecompressorReader( nspf, BlockHeader) pos = nspf.tell() if not useBlockCompression: decompressor = ZstdDecompressor().stream_reader(nspf) hash = sha256() f.write(header) hash.update(header) for s in sections: i = s.offset crypto = AESCTR(s.cryptoKey, s.cryptoCounter) end = s.offset + s.size while i < end: crypto.seek(i) chunkSz = 0x10000 if end - i > 0x10000 else end - i if useBlockCompression: inputChunk = blockDecompressorReader.read(chunkSz) else: inputChunk = decompressor.read(chunkSz) if not len(inputChunk): break if not useBlockCompression: decompressor.flush() if s.cryptoType in (3, 4): inputChunk = crypto.encrypt(inputChunk) f.write(inputChunk) hash.update(inputChunk) i += len(inputChunk) hexHash = hash.hexdigest() end = f.tell() written = (end - start) return (written, hexHash)
def decompress(self, fobj: IO[bytes]) -> IO[bytes]: decompressor = ZstdDecompressor() outfobj = NamedTemporaryFile(delete=False) try: decompressor.copy_stream(fobj, outfobj) outfobj.seek(0) yield outfobj finally: outfobj.close() remove(outfobj.name)
def open_read(self, path: str) -> IO[bytes]: decompressor = ZstdDecompressor() outfobj = NamedTemporaryFile(delete=False) try: with open(path, 'rb') as infobj: decompressor.copy_stream(infobj, outfobj) outfobj.seek(0) yield outfobj finally: outfobj.close() remove(outfobj.name)
def __decompressBlock(self, blockID): if (blockID >= len(self.CompressedBlockOffsetList)): raise EOFError( "BlockID exceeds the amounts of compressed blocks in that file!" ) self.nspf.seek(self.CompressedBlockOffsetList[blockID]) if self.CompressedBlockSizeList[blockID] < self.BlockSize: return ZstdDecompressor().decompress(self.nspf.read( self.BlockSize)) else: return self.nspf.read(self.BlockSize)
def decompress(data, compressor_id): if compressor_id == SnappyContext.compressor_id: # python-snappy doesn't support the buffer interface. # https://github.com/andrix/python-snappy/issues/65 # This only matters when data is a memoryview since # id(bytes(data)) == id(data) when data is a bytes. return snappy.uncompress(bytes(data)) elif compressor_id == ZlibContext.compressor_id: return zlib.decompress(data) elif compressor_id == ZstdContext.compressor_id: # ZstdDecompressor is not thread safe. # TODO: Use a pool? return ZstdDecompressor().decompress(data) else: raise ValueError("Unknown compressorId %d" % (compressor_id, ))
def _get_test_data(contest_id: str, problem_id: str, test_id: str, is_input: bool) -> Response: zctx = ZstdDecompressor() from io import BytesIO with transaction() as s: _ = _validate_token(s, admin_required=True) tc = s.query(TestCase).filter(TestCase.contest_id == contest_id, TestCase.problem_id == problem_id, TestCase.id == test_id).first() if not tc: abort(404) f = BytesIO(zctx.decompress(tc.input if is_input else tc.output)) return send_file(f, as_attachment=True, attachment_filename='{}.{}'.format( test_id, 'in' if is_input else 'out'))
def from_file(cls, file): data = [] # TODO: Unified "file loader" # (detect_codec is used is many place) codec = detect_compression(file) with open(file, "rb") as f: if codec == CompressionFormat.Zstandard: ctx = ZstdDecompressor() f = ctx.stream_reader(f) f = TextIOWrapper(f, "utf-8") for line in f: if line.startswith(";"): continue prefix, origins = line.split("\t") origins = [int(x) for x in origins.split(",")] data.append((prefix, origins)) return cls(data)
def __decompressBlock(self, blockID): if self.CurrentBlockId == blockID: return self.CurrentBlock decompressedBlockSize = self.BlockSize if blockID >= len(self.CompressedBlockOffsetList) - 1: if blockID >= len(self.CompressedBlockOffsetList): raise EOFError( "BlockID exceeds the amounts of compressed blocks in that file!" ) decompressedBlockSize = self.BlockHeader.decompressedSize % BlockSize self.nspf.seek(self.CompressedBlockOffsetList[blockID]) if self.CompressedBlockSizeList[blockID] < decompressedBlockSize: self.CurrentBlock = ZstdDecompressor().decompress( self.nspf.read(decompressedBlockSize)) else: self.CurrentBlock = self.nspf.read(decompressedBlockSize) self.CurrentBlockId = blockID return self.CurrentBlock
def __init__( self, path: Path, *, encoding: str, warn_uncompressed: bool = True, progress_bar: bool = False, progress_bar_desc: Optional[str] = None, ): self.path = path self._fp = path.open("rb") self._fin: BinaryIO if path.suffix == ".gz": self._fin = cast(BinaryIO, GzipFile(fileobj=self._fp)) elif path.suffix == ".bz2": self._fin = cast(BinaryIO, BZ2File(self._fp)) elif path.suffix == ".xz": self._fin = cast(BinaryIO, LZMAFile(self._fp)) elif path.suffix == ".zst": self._fin = cast(BinaryIO, ZstdDecompressor().stream_reader(self._fp)) else: if warn_uncompressed: # pragma: no cover _LOGGER.warning( "Could not detect compression type of file '{}' from its " "extension, treating as uncompressed file.", path, ) self._fin = self._fp self._progress_bar: Optional[tqdm[None]] = None if progress_bar: self._progress_bar = tqdm( desc=progress_bar_desc or self.path.name, total=self.size(), unit="B", unit_scale=True, unit_divisor=1024, dynamic_ncols=True, ) super().__init__(self._fin, encoding=encoding)
def decompress(data, compressor_id): if compressor_id == SnappyContext.compressor_id: # python-snappy doesn't support the buffer interface. # https://github.com/andrix/python-snappy/issues/65 # This only matters when data is a memoryview since # id(bytes(data)) == id(data) when data is a bytes. # NOTE: bytes(memoryview) returns the memoryview repr # in Python 2.7. The right thing to do in 2.7 is call # memoryview.tobytes(), but we currently only use # memoryview in Python 3.x. return snappy.uncompress(bytes(data)) elif compressor_id == ZlibContext.compressor_id: return zlib.decompress(data) elif compressor_id == ZstdContext.compressor_id: # ZstdDecompressor is not thread safe. # TODO: Use a pool? return ZstdDecompressor().decompress(data) else: raise ValueError("Unknown compressorId %d" % (compressor_id, ))
def get_submission(contest_id: str, submission_id: str) -> Response: params, _ = _validate_request() zctx = ZstdDecompressor() with transaction() as s: u = _validate_token(s) contest = s.query(Contest).filter(Contest.id == contest_id).first() if not (contest and contest.is_accessible(u)): abort(404) tmp = s.query(Submission, User.name).filter(Submission.contest_id == contest_id, Submission.id == submission_id, Submission.user_id == User.id).first() if not tmp: abort(404) submission, user_name = tmp if not submission.is_accessible(contest, u): abort(404) ret = submission.to_dict() ret['user_name'] = user_name ret['tests'] = [] for t_raw in s.query(JudgeResult).filter( JudgeResult.submission_id == submission_id).order_by( JudgeResult.status, JudgeResult.test_id): t = t_raw.to_dict() # 不要な情報を削除 t.pop('contest_id') t.pop('problem_id') t.pop('submission_id') t['id'] = t['test_id'] t.pop('test_id') if not (contest.is_finished() or (u and u['admin'])): # コンテスト中&非管理者の場合は # 実行時間とメモリ消費量を返却しない # (NULLの場合はto_dictで設定されないのでpopの引数にNoneを指定) t.pop('time', None) t.pop('memory', None) ret['tests'].append(t) ret['code'] = zctx.decompress(ret['code']).decode('utf-8') return jsonify(ret)
def _read_rel_zs_rows(filepath, chunk_size=8 * 1000 * 1000): from zstandard import ZstdDecompressor with open(filepath, "rb") as fh: ctx = ZstdDecompressor() with ctx.stream_reader(fh) as reader: over = False chunks = [] rows = [] while not over: have_row = False while not have_row: chunk = reader.read(chunk_size) if not chunk: over = True break if b"\n" in chunk: have_row = True chunks.append(chunk) (new_rows, semi_row) = _consume_rows(chunks) rows += new_rows chunks = [semi_row] return rows
def __decompressNcz(nspf, f, statusReportInfo, pleaseNoPrint): ncaHeaderSize = 0x4000 blockID = 0 nspf.seek(0) header = nspf.read(ncaHeaderSize) if f != None: start = f.tell() magic = nspf.read(8) if not magic == b'NCZSECTN': raise ValueError("No NCZSECTN found! Is this really a .ncz file?") sectionCount = nspf.readInt64() sections = [Header.Section(nspf) for _ in range(sectionCount)] nca_size = ncaHeaderSize for i in range(sectionCount): nca_size += sections[i].size pos = nspf.tell() blockMagic = nspf.read(8) nspf.seek(pos) useBlockCompression = blockMagic == b'NCZBLOCK' blockSize = -1 if useBlockCompression: BlockHeader = Header.Block(nspf) blockDecompressorReader = BlockDecompressorReader.BlockDecompressorReader( nspf, BlockHeader) pos = nspf.tell() if not useBlockCompression: decompressor = ZstdDecompressor().stream_reader(nspf) hash = sha256() if statusReportInfo == None: BAR_FMT = u'{desc}{desc_pad}{percentage:3.0f}%|{bar}| {count:{len_total}d}/{total:d} {unit} [{elapsed}<{eta}, {rate:.2f}{unit_pad}{unit}/s]' bar = enlighten.Counter(total=nca_size // 1048576, desc='Decompress', unit="MiB", color='red', bar_format=BAR_FMT) decompressedBytes = len(header) if f != None: f.write(header) if statusReportInfo != None: statusReport, id = statusReportInfo statusReport[id] = [len(header), 0, nca_size] else: bar.count = decompressedBytes // 1048576 bar.refresh() hash.update(header) for s in sections: i = s.offset crypto = aes128.AESCTR(s.cryptoKey, s.cryptoCounter) end = s.offset + s.size while i < end: crypto.seek(i) chunkSz = 0x10000 if end - i > 0x10000 else end - i if useBlockCompression: inputChunk = blockDecompressorReader.read(chunkSz) else: inputChunk = decompressor.read(chunkSz) if not len(inputChunk): break if not useBlockCompression: decompressor.flush() if s.cryptoType in (3, 4): inputChunk = crypto.encrypt(inputChunk) if f != None: f.write(inputChunk) hash.update(inputChunk) lenInputChunk = len(inputChunk) i += lenInputChunk decompressedBytes += lenInputChunk if statusReportInfo != None: statusReport[id] = [ statusReport[id][0] + chunkSz, statusReport[id][1], nca_size ] else: bar.count = decompressedBytes // 1048576 bar.refresh() bar.close() hexHash = hash.hexdigest() if f != None: end = f.tell() written = (end - start) return (written, hexHash) return (0, hexHash)
def __decompress(filePath, outputDir = None, write = True, raiseVerificationException = False): ncaHeaderSize = 0x4000 CHUNK_SZ = 0x100000 if write: nspPath = Path(filePath[0:-1] + 'p') if outputDir == None else Path(outputDir).joinpath(Path(filePath[0:-1] + 'p').name).resolve(strict=False) Print.info('decompressing %s -> %s' % (filePath, nspPath)) newNsp = Pfs0.Pfs0Stream(nspPath) fileHashes = FileExistingChecks.ExtractHashes(filePath) filePath = str(Path(filePath).resolve()) container = factory(filePath) container.open(filePath, 'rb') for nspf in container: if isinstance(nspf, Nca.Nca) and nspf.header.contentType == Type.Content.DATA: Print.info('skipping delta fragment') continue if not nspf._path.endswith('.ncz'): verifyFile = nspf._path.endswith('.nca') and not nspf._path.endswith('.cnmt.nca') if write: f = newNsp.add(nspf._path, nspf.size) hash = sha256() nspf.seek(0) while not nspf.eof(): inputChunk = nspf.read(CHUNK_SZ) hash.update(inputChunk) if write: f.write(inputChunk) if verifyFile: if hash.hexdigest() in fileHashes: Print.error('[VERIFIED] {0}'.format(nspf._path)) else: Print.info('[CORRUPTED] {0}'.format(nspf._path)) if raiseVerificationException: raise Exception("Verification detected hash missmatch!") elif not write: Print.info('[EXISTS] {0}'.format(nspf._path)) continue newFileName = nspf._path[0:-1] + 'a' if write: f = newNsp.add(newFileName, nspf.size) start = f.tell() blockID = 0 nspf.seek(0) header = nspf.read(ncaHeaderSize) magic = nspf.read(8) if not magic == b'NCZSECTN': raise ValueError("No NCZSECTN found! Is this really a .ncz file?") sectionCount = nspf.readInt64() sections = [Header.Section(nspf) for _ in range(sectionCount)] nca_size = ncaHeaderSize for i in range(sectionCount): nca_size += sections[i].size pos = nspf.tell() blockMagic = nspf.read(8) nspf.seek(pos) useBlockCompression = blockMagic == b'NCZBLOCK' blockSize = -1 if useBlockCompression: BlockHeader = Header.Block(nspf) blockDecompressorReader = BlockDecompressorReader.BlockDecompressorReader(nspf, BlockHeader) pos = nspf.tell() if not useBlockCompression: decompressor = ZstdDecompressor().stream_reader(nspf) hash = sha256() with tqdm(total=nca_size, unit_scale=True, unit="B") as bar: if write: f.write(header) bar.update(len(header)) hash.update(header) for s in sections: i = s.offset crypto = aes128.AESCTR(s.cryptoKey, s.cryptoCounter) end = s.offset + s.size while i < end: crypto.seek(i) chunkSz = 0x10000 if end - i > 0x10000 else end - i if useBlockCompression: inputChunk = blockDecompressorReader.read(chunkSz) else: inputChunk = decompressor.read(chunkSz) if not len(inputChunk): break if not useBlockCompression: decompressor.flush() if s.cryptoType in (3, 4): inputChunk = crypto.encrypt(inputChunk) if write: f.write(inputChunk) hash.update(inputChunk) i += len(inputChunk) bar.update(chunkSz) if hash.hexdigest() in fileHashes: Print.error('[VERIFIED] {0}'.format(nspf._path)) else: Print.info('[CORRUPTED] {0}'.format(nspf._path)) if raiseVerificationException: raise Exception("Verification detected hash missmatch") if write: end = f.tell() written = (end - start) newNsp.resize(newFileName, written) continue if write: newNsp.close() container.close()
def __decompressNcz(nspf, f, statusReportInfo, pleaseNoPrint): UNCOMPRESSABLE_HEADER_SIZE = 0x4000 blockID = 0 nspf.seek(0) header = nspf.read(UNCOMPRESSABLE_HEADER_SIZE) if f != None: start = f.tell() magic = nspf.read(8) if not magic == b'NCZSECTN': raise ValueError("No NCZSECTN found! Is this really a .ncz file?") sectionCount = nspf.readInt64() sections = [Header.Section(nspf) for _ in range(sectionCount)] if sections[0].offset - UNCOMPRESSABLE_HEADER_SIZE > 0: fakeSection = Header.FakeSection( UNCOMPRESSABLE_HEADER_SIZE, sections[0].offset - UNCOMPRESSABLE_HEADER_SIZE) sections.insert(0, fakeSection) nca_size = UNCOMPRESSABLE_HEADER_SIZE for i in range(sectionCount): nca_size += sections[i].size pos = nspf.tell() blockMagic = nspf.read(8) nspf.seek(pos) useBlockCompression = blockMagic == b'NCZBLOCK' blockSize = -1 if useBlockCompression: BlockHeader = Header.Block(nspf) blockDecompressorReader = BlockDecompressorReader.BlockDecompressorReader( nspf, BlockHeader) pos = nspf.tell() if not useBlockCompression: decompressor = ZstdDecompressor().stream_reader(nspf) hash = sha256() if statusReportInfo == None: BAR_FMT = u'{desc}{desc_pad}{percentage:3.0f}%|{bar}| {count:{len_total}d}/{total:d} {unit} [{elapsed}<{eta}, {rate:.2f}{unit_pad}{unit}/s]' bar = enlighten.Counter(total=nca_size // 1048576, desc='Decompress', unit="MiB", color='red', bar_format=BAR_FMT) decompressedBytes = len(header) if f != None: f.write(header) if statusReportInfo != None: statusReport, id = statusReportInfo statusReport[id] = [len(header), 0, nca_size] else: bar.count = decompressedBytes // 1048576 bar.refresh() hash.update(header) firstSection = True for s in sections: i = s.offset useCrypto = s.cryptoType in (3, 4) if useCrypto: crypto = aes128.AESCTR(s.cryptoKey, s.cryptoCounter) end = s.offset + s.size if firstSection: firstSection = False uncompressedSize = UNCOMPRESSABLE_HEADER_SIZE - sections[0].offset if uncompressedSize > 0: i += uncompressedSize while i < end: if useCrypto: crypto.seek(i) chunkSz = 0x10000 if end - i > 0x10000 else end - i if useBlockCompression: inputChunk = blockDecompressorReader.read(chunkSz) else: inputChunk = decompressor.read(chunkSz) decompressor.flush() if not len(inputChunk): break if useCrypto: inputChunk = crypto.encrypt(inputChunk) if f != None: f.write(inputChunk) hash.update(inputChunk) lenInputChunk = len(inputChunk) i += lenInputChunk decompressedBytes += lenInputChunk if statusReportInfo != None: statusReport[id] = [ statusReport[id][0] + chunkSz, statusReport[id][1], nca_size ] else: bar.count = decompressedBytes // 1048576 bar.refresh() if statusReportInfo == None: bar.close() #Line break after closing the process bar is required to prevent #the next output from being on the same line as the process bar print() hexHash = hash.hexdigest() if f != None: end = f.tell() written = (end - start) return (written, hexHash) return (0, hexHash)
def __init__(self): self.compressor = ZstdCompressor() self.decompressor = ZstdDecompressor()
def read_index(index_path: Path, rsa_priv_key_path: Path = None) -> dict: if index_path is None or not index_path.is_file(): raise RuntimeError( f"Unable to read non-existant index file \"{index_path}\"") encryption_flag = None compression_flag = None session_key = None data_size = None to_read_buffer = None with open(index_path, "rb") as index_stream: magic = str(index_stream.read(7)) if magic != "TINFOIL": raise RuntimeError( "Invalid tinfoil index magic.\n\nExpected Magic = " + f"\"TINFOIL\"\nMagic in index file = \"{magic}\"") flags = index_stream.read(1)[0] encryption_flag = flags & 0xF0 key_available = rsa_priv_key_path is not None and \ rsa_priv_key_path.is_file() if encryption_flag == EncryptionFlag.ENCRYPT and not key_available: raise RuntimeError( "Unable to decrypt encrypted index without private key.") compression_flag = flags & 0x0F if compression_flag not in CompressionFlag: raise RuntimeError( "Unimplemented compression method encountered while reading " + "index header.") session_key = index_stream.read(0x100) data_size = int.from_bytes(index_stream.read(8), byteorder="little") to_read_buffer = index_stream.read() if encryption_flag == EncryptionFlag.ENCRYPT: rsa_priv_key = import_rsa_key(open(rsa_priv_key_path).read()) pkcs1_oaep_ctx = new_pkcs1_oaep_ctx(rsa_priv_key, hashAlgo=SHA256, label=b"") aes_key = pkcs1_oaep_ctx.decrypt(session_key) aes_ctx = new_aes_ctx(aes_key, MODE_ECB) to_read_buffer = aes_ctx.decrypt(to_read_buffer) if compression_flag == CompressionFlag.ZSTD_COMPRESSION: to_read_buffer = ZstdDecompressor().decompress( to_read_buffer[:data_size]) elif compression_flag == CompressionFlag.ZLIB_COMPRESSION: to_read_buffer = zlib_decompress(to_read_buffer[:data_size]) elif compression_flag == CompressionFlag.NO_COMPRESSION: to_read_buffer = to_read_buffer[:data_size] try: return json_deserialize(to_read_buffer) except JSONDecodeError: raise RuntimeError("Unable to deserialize index data.")
def __init__(self, item_type: type, num_pages: int, page_size: int): self.item_type = item_type self.num_pages = num_pages self.page_size = page_size self.decompressor = ZstdDecompressor() self.mmap = None
def __decompressNcz(nspf, f, statusReportInfo): UNCOMPRESSABLE_HEADER_SIZE = 0x4000 blockID = 0 nspf.seek(0) header = nspf.read(UNCOMPRESSABLE_HEADER_SIZE) if f is not None: start = f.tell() magic = nspf.read(8) if not magic == b'NCZSECTN': raise ValueError("No NCZSECTN found! Is this really a .ncz file?") sectionCount = nspf.readInt64() sections = [Section(nspf) for _ in range(sectionCount)] if sections[0].offset - UNCOMPRESSABLE_HEADER_SIZE > 0: fakeSection = FakeSection( UNCOMPRESSABLE_HEADER_SIZE, sections[0].offset - UNCOMPRESSABLE_HEADER_SIZE) sections.insert(0, fakeSection) nca_size = UNCOMPRESSABLE_HEADER_SIZE for i in range(sectionCount): nca_size += sections[i].size decompressor = ZstdDecompressor().stream_reader(nspf) hash = sha256() bar = Status.create(nspf.size, desc=os.path.basename(nspf._path), unit='B') # if statusReportInfo == None: # BAR_FMT = u'{desc}{desc_pad}{percentage:3.0f}%|{bar}| {count:{len_total}d}/{total:d} {unit} [{elapsed}<{eta}, {rate:.2f}{unit_pad}{unit}/s]' # bar = enlighten.Counter(total=nca_size//1048576, desc='Decompress', unit="MiB", color='red', bar_format=BAR_FMT) decompressedBytes = len(header) if f is not None: f.write(header) bar.add(len(header)) hash.update(header) firstSection = True for s in sections: i = s.offset useCrypto = s.cryptoType in (3, 4) if useCrypto: crypto = aes128.AESCTR(s.cryptoKey, s.cryptoCounter) end = s.offset + s.size if firstSection: firstSection = False uncompressedSize = UNCOMPRESSABLE_HEADER_SIZE - sections[0].offset if uncompressedSize > 0: i += uncompressedSize while i < end: if useCrypto: crypto.seek(i) chunkSz = 0x10000 if end - i > 0x10000 else end - i inputChunk = decompressor.read(chunkSz) decompressor.flush() if not len(inputChunk): break if useCrypto: inputChunk = crypto.encrypt(inputChunk) if f is not None: f.write(inputChunk) bar.add(len(inputChunk)) hash.update(inputChunk) lenInputChunk = len(inputChunk) i += lenInputChunk decompressedBytes += lenInputChunk bar.add(lenInputChunk) bar.close() print() hexHash = hash.hexdigest() if f is not None: end = f.tell() written = (end - start) return (written, hexHash) return (0, hexHash)
else: if not sent: continue if para_id == '0': continue out += ' '.join(tokenizer.tokenize(sent)) + '\n' sent = '' if sent: out += ' '.join(tokenizer.tokenize(sent)) + '\n' return out dctx = ZstdDecompressor() for group in tqdm(os.listdir(source_dir), ncols=80, position=0): group_dir = os.path.join(source_dir, group) group_outs = [] for filename in tqdm(os.listdir(group_dir), ncols=80, desc=group, position=1): filepath = os.path.join(group_dir, filename) with open(filepath, 'rb') as f: with dctx.stream_reader(f) as r: text_stream = io.TextIOWrapper(r, encoding='utf-8') out = convert_file(text_stream) group_outs.append(out) out = '\n'.join(group_outs)