diff --git a/python/amunpy/src/amunpy/amunxml.py b/python/amunpy/src/amunpy/amunxml.py index 5bdbfda..78d4902 100644 --- a/python/amunpy/src/amunpy/amunxml.py +++ b/python/amunpy/src/amunpy/amunxml.py @@ -208,14 +208,18 @@ class AmunXML(Amun): return self.__swap__(dset) - def __check_digest(self, filename, digest, data): + def __check_digest(self, filename, hash_type, digest, data): ''' - Verifies if the provided digest matches the XXH64 hash of data - stored in the given filename. + Verifies if the provided digest matches the data. ''' import xxhash - if digest.lower() != xxhash.xxh64(data).hexdigest(): + failed = False + if hash_type == 'xxh64': + failed = digest.lower() != xxhash.xxh64(data).hexdigest() + elif hash_type == 'xxh3': + failed = digest.lower() != xxhash.xxh3_64(data).hexdigest() + if failed: print("File '{}' seems to be corrupted! Proceeding anyway...".format(filename)) @@ -235,7 +239,9 @@ class AmunXML(Amun): if 'compression_format' in self.binaries[dataset]: if 'compressed_digest' in self.binaries[dataset]: - self.__check_digest(fname, self.binaries[dataset]['compressed_digest'], stream) + htype = self.binaries[dataset]['digest_type'] + dhash = self.binaries[dataset]['compressed_digest'] + self.__check_digest(fname, htype, dhash, stream) comp = self.binaries[dataset]['compression_format'] if comp == 'zstd': @@ -248,12 +254,16 @@ class AmunXML(Amun): raise Exception("Binary file '{}' compressed in unsupported format {}!".format(fname, comp)) if 'digest' in self.binaries[dataset]: - self.__check_digest(fname, self.binaries[dataset]['digest'], data) + htype = self.binaries[dataset]['digest_type'] + dhash = self.binaries[dataset]['digest'] + self.__check_digest(fname, htype, dhash, data) return numpy.frombuffer(data, dtype=dtype) else: if 'digest' in self.binaries[dataset]: - self.__check_digest(fname, self.binaries[dataset]['digest'], stream) + htype = self.binaries[dataset]['digest_type'] + dhash = self.binaries[dataset]['digest'] + self.__check_digest(fname, htype, dhash, stream) return numpy.frombuffer(stream, dtype=dtype) else: @@ -276,7 +286,9 @@ class AmunXML(Amun): if 'compression_format' in self.chunks[chunk_number][dataset_name]: if 'compressed_digest' in self.chunks[chunk_number][dataset_name]: - self.__check_digest(fname, self.chunks[chunk_number][dataset_name]['compressed_digest'], stream) + htype = self.chunks[chunk_number][dataset_name]['digest_type'] + dhash = self.chunks[chunk_number][dataset_name]['compressed_digest'] + self.__check_digest(fname, htype, dhash, stream) comp = self.chunks[chunk_number][dataset_name]['compression_format'] if comp == 'zstd': @@ -290,12 +302,16 @@ class AmunXML(Amun): raise Exception("Binary file '{}' compressed in unsupported format {}!".format(fname, comp)) if 'digest' in self.chunks[chunk_number][dataset_name]: - self.__check_digest(fname, self.chunks[chunk_number][dataset_name]['digest'], data) + htype = self.chunks[chunk_number][dataset_name]['digest_type'] + dhash = self.chunks[chunk_number][dataset_name]['digest'] + self.__check_digest(fname, htype, dhash, data) return numpy.frombuffer(data, dtype=dtype) else: if 'digest' in self.chunks[chunk_number][dataset_name]: - self.__check_digest(fname, self.chunks[chunk_number][dataset_name]['digest'], stream) + htype = self.chunks[chunk_number][dataset_name]['digest_type'] + dhash = self.chunks[chunk_number][dataset_name]['digest'] + self.__check_digest(fname, htype, dhash, stream) return numpy.frombuffer(stream, dtype=dtype) else: