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*34917eadSPhilippe Mathieu-Daudéimport subprocess
16850a1951SThomas Huthimport tarfile
17850a1951SThomas Huth
18850a1951SThomas Huthdef archive_extract(archive, dest_dir, member=None):
19850a1951SThomas Huth    with tarfile.open(archive) as tf:
20850a1951SThomas Huth        if hasattr(tarfile, 'data_filter'):
21850a1951SThomas Huth            tf.extraction_filter = getattr(tarfile, 'data_filter',
22850a1951SThomas Huth                                           (lambda member, path: member))
23850a1951SThomas Huth        if member:
24850a1951SThomas Huth            tf.extract(member=member, path=dest_dir)
25850a1951SThomas Huth        else:
26850a1951SThomas Huth            tf.extractall(path=dest_dir)
27e2e9fd25SThomas Huth
28d5674412SThomas Huthdef gzip_uncompress(gz_path, output_path):
29d5674412SThomas Huth    if os.path.exists(output_path):
30d5674412SThomas Huth        return
31d5674412SThomas Huth    with gzip.open(gz_path, 'rb') as gz_in:
32d5674412SThomas Huth        try:
33d5674412SThomas Huth            with open(output_path, 'wb') as raw_out:
34d5674412SThomas Huth                shutil.copyfileobj(gz_in, raw_out)
35d5674412SThomas Huth        except:
36d5674412SThomas Huth            os.remove(output_path)
37d5674412SThomas Huth            raise
38d5674412SThomas Huth
39e2e9fd25SThomas Huthdef lzma_uncompress(xz_path, output_path):
40e2e9fd25SThomas Huth    if os.path.exists(output_path):
41e2e9fd25SThomas Huth        return
42e2e9fd25SThomas Huth    with lzma.open(xz_path, 'rb') as lzma_in:
43e2e9fd25SThomas Huth        try:
44e2e9fd25SThomas Huth            with open(output_path, 'wb') as raw_out:
45e2e9fd25SThomas Huth                shutil.copyfileobj(lzma_in, raw_out)
46e2e9fd25SThomas Huth        except:
47e2e9fd25SThomas Huth            os.remove(output_path)
48e2e9fd25SThomas Huth            raise
49*34917eadSPhilippe Mathieu-Daudé
50*34917eadSPhilippe Mathieu-Daudédef cpio_extract(cpio_handle, output_path):
51*34917eadSPhilippe Mathieu-Daudé    cwd = os.getcwd()
52*34917eadSPhilippe Mathieu-Daudé    os.chdir(output_path)
53*34917eadSPhilippe Mathieu-Daudé    subprocess.run(['cpio', '-i'],
54*34917eadSPhilippe Mathieu-Daudé                   input=cpio_handle.read(),
55*34917eadSPhilippe Mathieu-Daudé                   stderr=subprocess.DEVNULL)
56*34917eadSPhilippe Mathieu-Daudé    os.chdir(cwd)
57