1# Utilities for python-based QEMU tests
2#
3# Copyright 2024 Red Hat, Inc.
4#
5# Authors:
6#  Thomas Huth <thuth@redhat.com>
7#
8# This work is licensed under the terms of the GNU GPL, version 2 or
9# later.  See the COPYING file in the top-level directory.
10
11import gzip
12import lzma
13import os
14import shutil
15import subprocess
16import tarfile
17
18"""
19Round up to next power of 2
20"""
21def pow2ceil(x):
22    return 1 if x == 0 else 2**(x - 1).bit_length()
23
24def file_truncate(path, size):
25    if size != os.path.getsize(path):
26        with open(path, 'ab+') as fd:
27            fd.truncate(size)
28
29"""
30Expand file size to next power of 2
31"""
32def image_pow2ceil_expand(path):
33        size = os.path.getsize(path)
34        size_aligned = pow2ceil(size)
35        if size != size_aligned:
36            with open(path, 'ab+') as fd:
37                fd.truncate(size_aligned)
38
39def archive_extract(archive, dest_dir, member=None):
40    with tarfile.open(archive) as tf:
41        if hasattr(tarfile, 'data_filter'):
42            tf.extraction_filter = getattr(tarfile, 'data_filter',
43                                           (lambda member, path: member))
44        if member:
45            tf.extract(member=member, path=dest_dir)
46        else:
47            tf.extractall(path=dest_dir)
48
49def gzip_uncompress(gz_path, output_path):
50    if os.path.exists(output_path):
51        return
52    with gzip.open(gz_path, 'rb') as gz_in:
53        try:
54            with open(output_path, 'wb') as raw_out:
55                shutil.copyfileobj(gz_in, raw_out)
56        except:
57            os.remove(output_path)
58            raise
59
60def lzma_uncompress(xz_path, output_path):
61    if os.path.exists(output_path):
62        return
63    with lzma.open(xz_path, 'rb') as lzma_in:
64        try:
65            with open(output_path, 'wb') as raw_out:
66                shutil.copyfileobj(lzma_in, raw_out)
67        except:
68            os.remove(output_path)
69            raise
70
71def cpio_extract(cpio_handle, output_path):
72    cwd = os.getcwd()
73    os.chdir(output_path)
74    subprocess.run(['cpio', '-i'],
75                   input=cpio_handle.read(),
76                   stderr=subprocess.DEVNULL)
77    os.chdir(cwd)
78