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
11*d5674412SThomas Huthimport gzip
12e2e9fd25SThomas Huthimport lzma
13e2e9fd25SThomas Huthimport os
14e2e9fd25SThomas Huthimport shutil
15850a1951SThomas Huthimport tarfile
16850a1951SThomas Huth
17850a1951SThomas Huthdef archive_extract(archive, dest_dir, member=None):
18850a1951SThomas Huth    with tarfile.open(archive) as tf:
19850a1951SThomas Huth        if hasattr(tarfile, 'data_filter'):
20850a1951SThomas Huth            tf.extraction_filter = getattr(tarfile, 'data_filter',
21850a1951SThomas Huth                                           (lambda member, path: member))
22850a1951SThomas Huth        if member:
23850a1951SThomas Huth            tf.extract(member=member, path=dest_dir)
24850a1951SThomas Huth        else:
25850a1951SThomas Huth            tf.extractall(path=dest_dir)
26e2e9fd25SThomas Huth
27*d5674412SThomas Huthdef gzip_uncompress(gz_path, output_path):
28*d5674412SThomas Huth    if os.path.exists(output_path):
29*d5674412SThomas Huth        return
30*d5674412SThomas Huth    with gzip.open(gz_path, 'rb') as gz_in:
31*d5674412SThomas Huth        try:
32*d5674412SThomas Huth            with open(output_path, 'wb') as raw_out:
33*d5674412SThomas Huth                shutil.copyfileobj(gz_in, raw_out)
34*d5674412SThomas Huth        except:
35*d5674412SThomas Huth            os.remove(output_path)
36*d5674412SThomas Huth            raise
37*d5674412SThomas Huth
38e2e9fd25SThomas Huthdef lzma_uncompress(xz_path, output_path):
39e2e9fd25SThomas Huth    if os.path.exists(output_path):
40e2e9fd25SThomas Huth        return
41e2e9fd25SThomas Huth    with lzma.open(xz_path, 'rb') as lzma_in:
42e2e9fd25SThomas Huth        try:
43e2e9fd25SThomas Huth            with open(output_path, 'wb') as raw_out:
44e2e9fd25SThomas Huth                shutil.copyfileobj(lzma_in, raw_out)
45e2e9fd25SThomas Huth        except:
46e2e9fd25SThomas Huth            os.remove(output_path)
47e2e9fd25SThomas Huth            raise
48