| Simon Glass | 3ac7d83 | 2022-01-09 20:14:03 -0700 | [diff] [blame^] | 1 | # SPDX-License-Identifier: GPL-2.0+ |
| 2 | # Copyright 2022 Google LLC |
| 3 | # |
| 4 | """Utilities to compress and decompress data""" |
| 5 | |
| 6 | import struct |
| 7 | import tempfile |
| 8 | |
| 9 | from patman import tools |
| 10 | |
| 11 | def Compress(indata, algo, with_header=True): |
| 12 | """Compress some data using a given algorithm |
| 13 | |
| 14 | Note that for lzma this uses an old version of the algorithm, not that |
| 15 | provided by xz. |
| 16 | |
| 17 | This requires 'lz4' and 'lzma_alone' tools. It also requires an output |
| 18 | directory to be previously set up, by calling PrepareOutputDir(). |
| 19 | |
| 20 | Care is taken to use unique temporary files so that this function can be |
| 21 | called from multiple threads. |
| 22 | |
| 23 | Args: |
| 24 | indata: Input data to compress |
| 25 | algo: Algorithm to use ('none', 'gzip', 'lz4' or 'lzma') |
| 26 | |
| 27 | Returns: |
| 28 | Compressed data |
| 29 | """ |
| 30 | if algo == 'none': |
| 31 | return indata |
| 32 | fname = tempfile.NamedTemporaryFile(prefix='%s.comp.tmp' % algo, |
| 33 | dir=tools.GetOutputDir()).name |
| 34 | tools.WriteFile(fname, indata) |
| 35 | if algo == 'lz4': |
| 36 | data = tools.Run('lz4', '--no-frame-crc', '-B4', '-5', '-c', fname, |
| 37 | binary=True) |
| 38 | # cbfstool uses a very old version of lzma |
| 39 | elif algo == 'lzma': |
| 40 | outfname = tempfile.NamedTemporaryFile(prefix='%s.comp.otmp' % algo, |
| 41 | dir=tools.GetOutputDir()).name |
| 42 | tools.Run('lzma_alone', 'e', fname, outfname, '-lc1', '-lp0', '-pb0', |
| 43 | '-d8') |
| 44 | data = tools.ReadFile(outfname) |
| 45 | elif algo == 'gzip': |
| 46 | data = tools.Run('gzip', '-c', fname, binary=True) |
| 47 | else: |
| 48 | raise ValueError("Unknown algorithm '%s'" % algo) |
| 49 | if with_header: |
| 50 | hdr = struct.pack('<I', len(data)) |
| 51 | data = hdr + data |
| 52 | return data |
| 53 | |
| 54 | def Decompress(indata, algo, with_header=True): |
| 55 | """Decompress some data using a given algorithm |
| 56 | |
| 57 | Note that for lzma this uses an old version of the algorithm, not that |
| 58 | provided by xz. |
| 59 | |
| 60 | This requires 'lz4' and 'lzma_alone' tools. It also requires an output |
| 61 | directory to be previously set up, by calling PrepareOutputDir(). |
| 62 | |
| 63 | Args: |
| 64 | indata: Input data to decompress |
| 65 | algo: Algorithm to use ('none', 'gzip', 'lz4' or 'lzma') |
| 66 | |
| 67 | Returns: |
| 68 | Compressed data |
| 69 | """ |
| 70 | if algo == 'none': |
| 71 | return indata |
| 72 | if with_header: |
| 73 | data_len = struct.unpack('<I', indata[:4])[0] |
| 74 | indata = indata[4:4 + data_len] |
| 75 | fname = tools.GetOutputFilename('%s.decomp.tmp' % algo) |
| 76 | with open(fname, 'wb') as fd: |
| 77 | fd.write(indata) |
| 78 | if algo == 'lz4': |
| 79 | data = tools.Run('lz4', '-dc', fname, binary=True) |
| 80 | elif algo == 'lzma': |
| 81 | outfname = tools.GetOutputFilename('%s.decomp.otmp' % algo) |
| 82 | tools.Run('lzma_alone', 'd', fname, outfname) |
| 83 | data = tools.ReadFile(outfname, binary=True) |
| 84 | elif algo == 'gzip': |
| 85 | data = tools.Run('gzip', '-cd', fname, binary=True) |
| 86 | else: |
| 87 | raise ValueError("Unknown algorithm '%s'" % algo) |
| 88 | return data |