1850a1951SThomas Huth# Utilities for python-based QEMU tests 2850a1951SThomas Huth# 3850a1951SThomas Huth# Copyright 2024 Red Hat, Inc. 4850a1951SThomas Huth# 5850a1951SThomas Huth# Authors: 6850a1951SThomas Huth# Thomas Huth <thuth@redhat.com> 7850a1951SThomas Huth# 8850a1951SThomas Huth# This work is licensed under the terms of the GNU GPL, version 2 or 9850a1951SThomas Huth# later. See the COPYING file in the top-level directory. 10850a1951SThomas Huth 11d5674412SThomas Huthimport gzip 12e2e9fd25SThomas Huthimport lzma 13e2e9fd25SThomas Huthimport os 14e2e9fd25SThomas Huthimport shutil 15*cfcb4484SDaniel P. Berrangé 16*cfcb4484SDaniel P. Berrangéfrom .archive import tar_extract as archive_extract 17*cfcb4484SDaniel P. Berrangéfrom .archive import cpio_extract 18850a1951SThomas Huth 19f7d6b772SThomas Huth""" 20f7d6b772SThomas HuthRound up to next power of 2 21f7d6b772SThomas Huth""" 22f7d6b772SThomas Huthdef pow2ceil(x): 23f7d6b772SThomas Huth return 1 if x == 0 else 2**(x - 1).bit_length() 24f7d6b772SThomas Huth 25f7d6b772SThomas Huthdef file_truncate(path, size): 26f7d6b772SThomas Huth if size != os.path.getsize(path): 27f7d6b772SThomas Huth with open(path, 'ab+') as fd: 28f7d6b772SThomas Huth fd.truncate(size) 29f7d6b772SThomas Huth 30f7d6b772SThomas Huth""" 31f7d6b772SThomas HuthExpand file size to next power of 2 32f7d6b772SThomas Huth""" 33f7d6b772SThomas Huthdef image_pow2ceil_expand(path): 34f7d6b772SThomas Huth size = os.path.getsize(path) 35f7d6b772SThomas Huth size_aligned = pow2ceil(size) 36f7d6b772SThomas Huth if size != size_aligned: 37f7d6b772SThomas Huth with open(path, 'ab+') as fd: 38f7d6b772SThomas Huth fd.truncate(size_aligned) 39f7d6b772SThomas Huth 40d5674412SThomas Huthdef gzip_uncompress(gz_path, output_path): 41d5674412SThomas Huth if os.path.exists(output_path): 42d5674412SThomas Huth return 43d5674412SThomas Huth with gzip.open(gz_path, 'rb') as gz_in: 44d5674412SThomas Huth try: 45d5674412SThomas Huth with open(output_path, 'wb') as raw_out: 46d5674412SThomas Huth shutil.copyfileobj(gz_in, raw_out) 47d5674412SThomas Huth except: 48d5674412SThomas Huth os.remove(output_path) 49d5674412SThomas Huth raise 50d5674412SThomas Huth 51e2e9fd25SThomas Huthdef lzma_uncompress(xz_path, output_path): 52e2e9fd25SThomas Huth if os.path.exists(output_path): 53e2e9fd25SThomas Huth return 54e2e9fd25SThomas Huth with lzma.open(xz_path, 'rb') as lzma_in: 55e2e9fd25SThomas Huth try: 56e2e9fd25SThomas Huth with open(output_path, 'wb') as raw_out: 57e2e9fd25SThomas Huth shutil.copyfileobj(lzma_in, raw_out) 58e2e9fd25SThomas Huth except: 59e2e9fd25SThomas Huth os.remove(output_path) 60e2e9fd25SThomas Huth raise 61