| 1 | # SPDX-License-Identifier: GPL-2.0-or-later |
| 2 | # |
| 3 | # Utilities for python-based QEMU tests |
| 4 | # |
| 5 | # Copyright 2024 Red Hat, Inc. |
| 6 | # |
| 7 | # Authors: |
| 8 | # Thomas Huth <thuth@redhat.com> |
| 9 | |
| 10 | import gzip |
| 11 | import lzma |
| 12 | import os |
| 13 | import stat |
| 14 | import shutil |
| 15 | from urllib.parse import urlparse |
| 16 | from subprocess import run, CalledProcessError |
| 17 | |
| 18 | from .asset import Asset |
| 19 | |
| 20 | |
| 21 | def gzip_uncompress(gz_path, output_path): |
| 22 | if os.path.exists(output_path): |
| 23 | return |
| 24 | with gzip.open(gz_path, 'rb') as gz_in: |
| 25 | try: |
| 26 | with open(output_path, 'wb') as raw_out: |
| 27 | shutil.copyfileobj(gz_in, raw_out) |
| 28 | except: |
| 29 | os.remove(output_path) |
| 30 | raise |
| 31 | |
| 32 | def lzma_uncompress(xz_path, output_path): |
| 33 | if os.path.exists(output_path): |
| 34 | return |
| 35 | with lzma.open(xz_path, 'rb') as lzma_in: |
| 36 | try: |
| 37 | with open(output_path, 'wb') as raw_out: |
| 38 | shutil.copyfileobj(lzma_in, raw_out) |
| 39 | except: |
| 40 | os.remove(output_path) |
| 41 | raise |
| 42 | |
| 43 | |
| 44 | def zstd_uncompress(zstd_path, output_path): |
| 45 | if os.path.exists(output_path): |
| 46 | return |
| 47 | |
| 48 | try: |
| 49 | run(['zstd', "-f", "-d", zstd_path, |
| 50 | "-o", output_path], capture_output=True, check=True) |
| 51 | except CalledProcessError as e: |
| 52 | os.remove(output_path) |
| 53 | raise Exception( |
| 54 | f"Unable to decompress zstd file {zstd_path} with {e}") from e |
| 55 | |
| 56 | # zstd copies source archive permissions for the output |
| 57 | # file, so must make this writable for QEMU |
| 58 | os.chmod(output_path, stat.S_IRUSR | stat.S_IWUSR) |
| 59 | |
| 60 | |
| 61 | def uncompress(compressed, uncompressed, format=None): |
| 62 | ''' |
| 63 | @params compressed: filename, Asset, or file-like object to uncompress |
| 64 | @params uncompressed: filename to uncompress into |
| 65 | @params format: optional compression format (gzip, lzma) |
| 66 | |
| 67 | Uncompresses @compressed into @uncompressed |
| 68 | |
| 69 | If @format is None, heuristics will be applied to guess the |
| 70 | format from the filename or Asset URL. @format must be non-None |
| 71 | if @uncompressed is a file-like object. |
| 72 | |
| 73 | Returns the fully qualified path to the uncompessed file |
| 74 | ''' |
| 75 | if format is None: |
| 76 | format = guess_uncompress_format(compressed) |
| 77 | |
| 78 | if format == "xz": |
| 79 | lzma_uncompress(str(compressed), uncompressed) |
| 80 | elif format == "gz": |
| 81 | gzip_uncompress(str(compressed), uncompressed) |
| 82 | elif format == "zstd": |
| 83 | zstd_uncompress(str(compressed), uncompressed) |
| 84 | else: |
| 85 | raise Exception(f"Unknown compression format {format}") |
| 86 | |
| 87 | def guess_uncompress_format(compressed): |
| 88 | ''' |
| 89 | @params compressed: filename, Asset, or file-like object to guess |
| 90 | |
| 91 | Guess the format of @compressed, raising an exception if |
| 92 | no format can be determined |
| 93 | ''' |
| 94 | if isinstance(compressed, Asset): |
| 95 | compressed = urlparse(compressed.url).path |
| 96 | elif not isinstance(compressed, str): |
| 97 | raise Exception(f"Unable to guess compression cformat for {compressed}") |
| 98 | |
| 99 | (_name, ext) = os.path.splitext(compressed) |
| 100 | if ext == ".xz": |
| 101 | return "xz" |
| 102 | elif ext == ".gz": |
| 103 | return "gz" |
| 104 | elif ext in [".zstd", ".zst"]: |
| 105 | return 'zstd' |
| 106 | else: |
| 107 | raise Exception(f"Unknown compression format for {compressed}") |