Пример #1
0
def __decompressNsz(filePath, outputDir, write, raiseVerificationException, statusReportInfo, pleaseNoPrint):
	fileHashes = FileExistingChecks.ExtractHashes(filePath)
	container = factory(filePath)
	container.open(str(filePath), 'rb')
	
	if write:
		filename = changeExtension(filePath, '.nsp')
		outPath = filename if outputDir == None else str(Path(outputDir).joinpath(filename))
		Print.info('Decompressing %s -> %s' % (filePath, outPath), pleaseNoPrint)
		with Pfs0.Pfs0Stream(outPath) as nsp:
			__decompressContainer(container, nsp, fileHashes, write, raiseVerificationException, statusReportInfo, pleaseNoPrint)
	else:
		__decompressContainer(container, None, fileHashes, write, raiseVerificationException, statusReportInfo, pleaseNoPrint)

	container.close()
Пример #2
0
def blockCompressNsp(filePath, compressionLevel, blockSizeExponent, outputDir,
                     threads):
    filePath = filePath.resolve()
    container = factory(filePath)
    container.open(str(filePath), 'rb')
    nszPath = outputDir.joinpath(filePath.stem + '.nsz')

    Print.info('Block compressing (level {0}) {1} -> {2}'.format(
        compressionLevel, filePath, nszPath))

    try:
        with Pfs0.Pfs0Stream(str(nszPath)) as nsp:
            blockCompressContainer(container, nsp, compressionLevel,
                                   blockSizeExponent, threads)
    except BaseException as ex:
        if not ex is KeyboardInterrupt:
            Print.error(format_exc())
        if nszPath.is_file():
            nszPath.unlink()

    container.close()
    return nszPath
Пример #3
0
def solidCompressNsp(filePath, compressionLevel, outputDir, threads,
                     stusReport, id, pleaseNoPrint):
    filePath = filePath.resolve()
    container = factory(filePath)
    container.open(str(filePath), 'rb')
    nszPath = outputDir.joinpath(filePath.stem + '.nsz')

    Print.info(
        'Solid compressing (level {0}) {1} -> {2}'.format(
            compressionLevel, filePath, nszPath), pleaseNoPrint)

    try:
        with Pfs0.Pfs0Stream(str(nszPath)) as nsp:
            processContainer(container, nsp, compressionLevel, threads,
                             stusReport, id, pleaseNoPrint)
    except BaseException as ex:
        if not ex is KeyboardInterrupt:
            Print.error(format_exc())
        if nszPath.is_file():
            nszPath.unlink()

    container.close()
    return nszPath
Пример #4
0
def blockCompress(filePath,
                  compressionLevel=18,
                  blockSizeExponent=20,
                  outputDir=None,
                  threads=32):
    CHUNK_SZ = 0x100000
    ncaHeaderSize = 0x4000
    if blockSizeExponent < 14 or blockSizeExponent > 32:
        raise ValueError("Block size must be between 14 and 32")
    blockSize = 2**blockSizeExponent
    filePath = str(Path(filePath).resolve(strict=False))
    container = factory(filePath)
    container.open(filePath, 'rb')
    nszPath = str(
        Path(filePath[0:-1] + 'z' if outputDir == None else str(
            Path(outputDir).joinpath(Path(filePath[0:-1] +
                                          'z').name))).resolve(strict=False))
    Print.info('compressing (level %d) %s -> %s' %
               (compressionLevel, filePath, nszPath))
    newNsp = Pfs0.Pfs0Stream(nszPath)
    try:
        manager = Manager()
        results = manager.list()
        readyForWork = Counter(0)
        pleaseKillYourself = Counter(0)
        TasksPerChunk = 209715200 // blockSize
        for i in range(TasksPerChunk):
            results.append(b"")
        work = manager.Queue(threads)
        pool = []

        for i in range(threads):
            p = Process(target=compressBlockTask,
                        args=(work, results, readyForWork, pleaseKillYourself))
            p.start()
            pool.append(p)

        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')
                    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()
                    bytesToCompress = nspf.size - ncaHeaderSize
                    blocksToCompress = bytesToCompress // blockSize + (
                        bytesToCompress % blockSize > 0)
                    compressedblockSizeList = [0] * blocksToCompress
                    header = b'NCZBLOCK'  #Magic
                    header += b'\x02'  #Version
                    header += b'\x01'  #Type
                    header += b'\x00'  #Unused
                    header += blockSizeExponent.to_bytes(
                        1, 'little')  #blockSizeExponent in bits: 2^x
                    header += blocksToCompress.to_bytes(
                        4, 'little')  #Amount of Blocks
                    header += bytesToCompress.to_bytes(
                        8, 'little')  #Decompressed Size
                    header += b'\x00' * (blocksToCompress * 4)
                    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())
                        while True:
                            buffer = partitions[partNr].read(blockSize)
                            while (len(buffer) < blockSize
                                   and partNr < len(partitions) - 1):
                                partitions[partNr].close()
                                partitions[partNr] = None
                                partNr += 1
                                buffer += partitions[partNr].read(blockSize -
                                                                  len(buffer))
                            if chunkRelativeBlockID >= TasksPerChunk or len(
                                    buffer) == 0:
                                while readyForWork.value() < threads:
                                    sleep(0.02)

                                for i in range(
                                        min(
                                            TasksPerChunk, blocksToCompress -
                                            startChunkBlockID)):
                                    compressedblockSizeList[startChunkBlockID +
                                                            i] = len(
                                                                results[i])
                                    f.write(results[i])
                                    results[i] = b""

                                if len(buffer) == 0:
                                    pleaseKillYourself.increment()

                                    for i in range(readyForWork.value()):
                                        work.put(None)

                                    while readyForWork.value() > 0:
                                        sleep(0.02)
                                    break
                                chunkRelativeBlockID = 0
                                startChunkBlockID = blockID
                            work.put([
                                buffer, compressionLevel,
                                compressedblockSizeList, chunkRelativeBlockID
                            ])
                            blockID += 1
                            chunkRelativeBlockID += 1
                            decompressedBytes += len(buffer)
                            bar.update(len(buffer))
                    partitions[partNr].close()
                    partitions[partNr] = None
                    f.seek(blocksHeaderFilePos + 24)
                    header = b""

                    for compressedblockSize in compressedblockSizeList:
                        header += compressedblockSize.to_bytes(4, 'little')

                    f.write(header)
                    f.seek(0, 2)  #Seek to end of file.
                    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
Пример #5
0
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
Пример #6
0
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()