def open_write(self, path: str) -> IO[bytes]: compressor = ZstdCompressor(level=self.level) fobj = open(path, 'wb') try: with compressor.stream_writer(fobj) as writer: yield writer finally: fobj.close()
def processContainer(readContainer, writeContainer, compressionLevel, threads, stusReport, id, pleaseNoPrint): for nspf in readContainer: if isinstance( nspf, Nca.Nca) and nspf.header.contentType == Type.Content.DATA: Print.info('Skipping delta fragment {0}'.format(nspf._path)) continue if isinstance(nspf, Nca.Nca) and ( nspf.header.contentType == Type.Content.PROGRAM or nspf.header.contentType == Type.Content.PUBLICDATA ) and nspf.size > UNCOMPRESSABLE_HEADER_SIZE: if isNcaPacked(nspf): offsetFirstSection = sortedFs(nspf)[0].offset newFileName = nspf._path[0:-1] + 'z' with writeContainer.add(newFileName, nspf.size, pleaseNoPrint) as f: start = f.tell() nspf.seek(0) f.write(nspf.read(UNCOMPRESSABLE_HEADER_SIZE)) sections = [] for fs in sortedFs(nspf): sections += fs.getEncryptionSections() if len(sections) == 0: raise Exception( "NCA can't be decrypted. Outdated keys.txt?") header = b'NCZSECTN' header += len(sections).to_bytes(8, 'little') i = 0 for fs in sections: i += 1 header += fs.offset.to_bytes(8, 'little') header += fs.size.to_bytes(8, 'little') header += fs.cryptoType.to_bytes(8, 'little') header += b'\x00' * 8 header += fs.cryptoKey header += fs.cryptoCounter f.write(header) blockID = 0 chunkRelativeBlockID = 0 startChunkBlockID = 0 blocksHeaderFilePos = f.tell() compressedblockSizeList = [] decompressedBytes = UNCOMPRESSABLE_HEADER_SIZE stusReport[id] = [0, 0, nspf.size] partitions = [] if offsetFirstSection - UNCOMPRESSABLE_HEADER_SIZE > 0: partitions.append( nspf.partition(offset=UNCOMPRESSABLE_HEADER_SIZE, size=offsetFirstSection - UNCOMPRESSABLE_HEADER_SIZE, cryptoType=Type.Crypto.CTR.NONE, autoOpen=True)) for section in sections: #Print.info('offset: %x\t\tsize: %x\t\ttype: %d\t\tiv%s' % (section.offset, section.size, section.cryptoType, str(hx(section.cryptoCounter))), pleaseNoPrint) partitions.append( nspf.partition(offset=section.offset, size=section.size, cryptoType=section.cryptoType, cryptoKey=section.cryptoKey, cryptoCounter=bytearray( section.cryptoCounter), autoOpen=True)) if UNCOMPRESSABLE_HEADER_SIZE - offsetFirstSection > 0: partitions[0].seek(UNCOMPRESSABLE_HEADER_SIZE - offsetFirstSection) partNr = 0 stusReport[id] = [nspf.tell(), f.tell(), nspf.size] if threads > 1: cctx = ZstdCompressor(level=compressionLevel, threads=threads) else: cctx = ZstdCompressor(level=compressionLevel) compressor = cctx.stream_writer(f) while True: buffer = partitions[partNr].read(CHUNK_SZ) while (len(buffer) < CHUNK_SZ and partNr < len(partitions) - 1): partitions[partNr].close() partitions[partNr] = None partNr += 1 buffer += partitions[partNr].read(CHUNK_SZ - len(buffer)) if len(buffer) == 0: break compressor.write(buffer) decompressedBytes += len(buffer) stusReport[id] = [nspf.tell(), f.tell(), nspf.size] partitions[partNr].close() partitions[partNr] = None compressor.flush(FLUSH_FRAME) compressor.flush(COMPRESSOBJ_FLUSH_FINISH) stusReport[id] = [nspf.tell(), f.tell(), nspf.size] written = f.tell() - start Print.info( 'Compressed {0}% {1} -> {2} - {3}'.format( written * 100 / nspf.size, decompressedBytes, written, nspf._path), pleaseNoPrint) writeContainer.resize(newFileName, written) continue else: Print.info('Skipping not packed {0}'.format(nspf._path)) with writeContainer.add(nspf._path, nspf.size, pleaseNoPrint) as f: nspf.seek(0) while not nspf.eof(): buffer = nspf.read(CHUNK_SZ) f.write(buffer)
def compress(self, fobj: IO[bytes]) -> IO[bytes]: compressor = ZstdCompressor(level=self.level) with compressor.stream_writer(fobj) as writer: yield writer
def solidCompress(filePath, compressionLevel=18, outputDir=None, threads=-1): ncaHeaderSize = 0x4000 filePath = str(Path(filePath).resolve()) container = factory(filePath) container.open(filePath, 'rb') CHUNK_SZ = 0x1000000 nszPath = str( Path(filePath[0:-1] + 'z' if outputDir == None else Path(outputDir). joinpath(Path(filePath[0:-1] + 'z').name)).resolve(strict=False)) for nspf in container: if isinstance(nspf, Ticket.Ticket): nspf.getRightsId() break # No need to go for other objects Print.info('compressing (level %d) %s -> %s' % (compressionLevel, filePath, nszPath)) newNsp = Pfs0.Pfs0Stream(nszPath) try: for nspf in container: if isinstance( nspf, Nca.Nca) and nspf.header.contentType == Type.Content.DATA: Print.info('skipping delta fragment') continue if isinstance(nspf, Nca.Nca) and ( nspf.header.contentType == Type.Content.PROGRAM or nspf.header.contentType == Type.Content.PUBLICDATA): if isNcaPacked(nspf, ncaHeaderSize): newFileName = nspf._path[0:-1] + 'z' f = newNsp.add(newFileName, nspf.size) start = f.tell() nspf.seek(0) f.write(nspf.read(ncaHeaderSize)) sections = [] for fs in sortedFs(nspf): sections += fs.getEncryptionSections() if len(sections) == 0: raise Exception( "NCA can't be decrypted. Outdated keys.txt?") header = b'NCZSECTN' header += len(sections).to_bytes(8, 'little') for fs in sections: header += fs.offset.to_bytes(8, 'little') header += fs.size.to_bytes(8, 'little') header += fs.cryptoType.to_bytes(8, 'little') header += b'\x00' * 8 header += fs.cryptoKey header += fs.cryptoCounter f.write(header) decompressedBytes = ncaHeaderSize with tqdm(total=nspf.size, unit_scale=True, unit="B") as bar: partitions = [ nspf.partition( offset=section.offset, size=section.size, n=None, cryptoType=section.cryptoType, cryptoKey=section.cryptoKey, cryptoCounter=bytearray(section.cryptoCounter), autoOpen=True) for section in sections ] partNr = 0 bar.update(f.tell()) cctx = ZstdCompressor( level=compressionLevel, threads=threads ) if threads > 1 else ZstdCompressor( level=compressionLevel) compressor = cctx.stream_writer(f) while True: buffer = partitions[partNr].read(CHUNK_SZ) while (len(buffer) < CHUNK_SZ and partNr < len(partitions) - 1): partitions[partNr].close() partitions[partNr] = None partNr += 1 buffer += partitions[partNr].read(CHUNK_SZ - len(buffer)) if len(buffer) == 0: break compressor.write(buffer) decompressedBytes += len(buffer) bar.update(len(buffer)) partitions[partNr].close() partitions[partNr] = None compressor.flush(FLUSH_FRAME) compressor.flush(COMPRESSOBJ_FLUSH_FINISH) written = f.tell() - start print('compressed %d%% %d -> %d - %s' % (int(written * 100 / nspf.size), decompressedBytes, written, nspf._path)) newNsp.resize(newFileName, written) continue else: print('not packed!') f = newNsp.add(nspf._path, nspf.size) nspf.seek(0) while not nspf.eof(): buffer = nspf.read(CHUNK_SZ) f.write(buffer) except KeyboardInterrupt: remove(nszPath) raise KeyboardInterrupt except BaseException: Print.error(format_exc()) remove(nszPath) finally: newNsp.close() container.close() return nszPath