patman: Drop references to __future__
[oweals/u-boot.git] / tools / patman / tools.py
1 # SPDX-License-Identifier: GPL-2.0+
2 #
3 # Copyright (c) 2016 Google, Inc
4 #
5
6 import command
7 import glob
8 import os
9 import shutil
10 import struct
11 import sys
12 import tempfile
13
14 import tout
15
16 # Output directly (generally this is temporary)
17 outdir = None
18
19 # True to keep the output directory around after exiting
20 preserve_outdir = False
21
22 # Path to the Chrome OS chroot, if we know it
23 chroot_path = None
24
25 # Search paths to use for Filename(), used to find files
26 search_paths = []
27
28 tool_search_paths = []
29
30 # Tools and the packages that contain them, on debian
31 packages = {
32     'lz4': 'liblz4-tool',
33     }
34
35 # List of paths to use when looking for an input file
36 indir = []
37
38 def PrepareOutputDir(dirname, preserve=False):
39     """Select an output directory, ensuring it exists.
40
41     This either creates a temporary directory or checks that the one supplied
42     by the user is valid. For a temporary directory, it makes a note to
43     remove it later if required.
44
45     Args:
46         dirname: a string, name of the output directory to use to store
47                 intermediate and output files. If is None - create a temporary
48                 directory.
49         preserve: a Boolean. If outdir above is None and preserve is False, the
50                 created temporary directory will be destroyed on exit.
51
52     Raises:
53         OSError: If it cannot create the output directory.
54     """
55     global outdir, preserve_outdir
56
57     preserve_outdir = dirname or preserve
58     if dirname:
59         outdir = dirname
60         if not os.path.isdir(outdir):
61             try:
62                 os.makedirs(outdir)
63             except OSError as err:
64                 raise CmdError("Cannot make output directory '%s': '%s'" %
65                                 (outdir, err.strerror))
66         tout.Debug("Using output directory '%s'" % outdir)
67     else:
68         outdir = tempfile.mkdtemp(prefix='binman.')
69         tout.Debug("Using temporary directory '%s'" % outdir)
70
71 def _RemoveOutputDir():
72     global outdir
73
74     shutil.rmtree(outdir)
75     tout.Debug("Deleted temporary directory '%s'" % outdir)
76     outdir = None
77
78 def FinaliseOutputDir():
79     global outdir, preserve_outdir
80
81     """Tidy up: delete output directory if temporary and not preserved."""
82     if outdir and not preserve_outdir:
83         _RemoveOutputDir()
84         outdir = None
85
86 def GetOutputFilename(fname):
87     """Return a filename within the output directory.
88
89     Args:
90         fname: Filename to use for new file
91
92     Returns:
93         The full path of the filename, within the output directory
94     """
95     return os.path.join(outdir, fname)
96
97 def _FinaliseForTest():
98     """Remove the output directory (for use by tests)"""
99     global outdir
100
101     if outdir:
102         _RemoveOutputDir()
103         outdir = None
104
105 def SetInputDirs(dirname):
106     """Add a list of input directories, where input files are kept.
107
108     Args:
109         dirname: a list of paths to input directories to use for obtaining
110                 files needed by binman to place in the image.
111     """
112     global indir
113
114     indir = dirname
115     tout.Debug("Using input directories %s" % indir)
116
117 def GetInputFilename(fname):
118     """Return a filename for use as input.
119
120     Args:
121         fname: Filename to use for new file
122
123     Returns:
124         The full path of the filename, within the input directory
125     """
126     if not indir or fname[:1] == '/':
127         return fname
128     for dirname in indir:
129         pathname = os.path.join(dirname, fname)
130         if os.path.exists(pathname):
131             return pathname
132
133     raise ValueError("Filename '%s' not found in input path (%s) (cwd='%s')" %
134                      (fname, ','.join(indir), os.getcwd()))
135
136 def GetInputFilenameGlob(pattern):
137     """Return a list of filenames for use as input.
138
139     Args:
140         pattern: Filename pattern to search for
141
142     Returns:
143         A list of matching files in all input directories
144     """
145     if not indir:
146         return glob.glob(fname)
147     files = []
148     for dirname in indir:
149         pathname = os.path.join(dirname, pattern)
150         files += glob.glob(pathname)
151     return sorted(files)
152
153 def Align(pos, align):
154     if align:
155         mask = align - 1
156         pos = (pos + mask) & ~mask
157     return pos
158
159 def NotPowerOfTwo(num):
160     return num and (num & (num - 1))
161
162 def SetToolPaths(toolpaths):
163     """Set the path to search for tools
164
165     Args:
166         toolpaths: List of paths to search for tools executed by Run()
167     """
168     global tool_search_paths
169
170     tool_search_paths = toolpaths
171
172 def PathHasFile(path_spec, fname):
173     """Check if a given filename is in the PATH
174
175     Args:
176         path_spec: Value of PATH variable to check
177         fname: Filename to check
178
179     Returns:
180         True if found, False if not
181     """
182     for dir in path_spec.split(':'):
183         if os.path.exists(os.path.join(dir, fname)):
184             return True
185     return False
186
187 def Run(name, *args, **kwargs):
188     """Run a tool with some arguments
189
190     This runs a 'tool', which is a program used by binman to process files and
191     perhaps produce some output. Tools can be located on the PATH or in a
192     search path.
193
194     Args:
195         name: Command name to run
196         args: Arguments to the tool
197
198     Returns:
199         CommandResult object
200     """
201     try:
202         binary = kwargs.get('binary')
203         env = None
204         if tool_search_paths:
205             env = dict(os.environ)
206             env['PATH'] = ':'.join(tool_search_paths) + ':' + env['PATH']
207         all_args = (name,) + args
208         result = command.RunPipe([all_args], capture=True, capture_stderr=True,
209                                  env=env, raise_on_error=False, binary=binary)
210         if result.return_code:
211             raise Exception("Error %d running '%s': %s" %
212                (result.return_code,' '.join(all_args),
213                 result.stderr))
214         return result.stdout
215     except:
216         if env and not PathHasFile(env['PATH'], name):
217             msg = "Please install tool '%s'" % name
218             package = packages.get(name)
219             if package:
220                  msg += " (e.g. from package '%s')" % package
221             raise ValueError(msg)
222         raise
223
224 def Filename(fname):
225     """Resolve a file path to an absolute path.
226
227     If fname starts with ##/ and chroot is available, ##/ gets replaced with
228     the chroot path. If chroot is not available, this file name can not be
229     resolved, `None' is returned.
230
231     If fname is not prepended with the above prefix, and is not an existing
232     file, the actual file name is retrieved from the passed in string and the
233     search_paths directories (if any) are searched to for the file. If found -
234     the path to the found file is returned, `None' is returned otherwise.
235
236     Args:
237       fname: a string,  the path to resolve.
238
239     Returns:
240       Absolute path to the file or None if not found.
241     """
242     if fname.startswith('##/'):
243       if chroot_path:
244         fname = os.path.join(chroot_path, fname[3:])
245       else:
246         return None
247
248     # Search for a pathname that exists, and return it if found
249     if fname and not os.path.exists(fname):
250         for path in search_paths:
251             pathname = os.path.join(path, os.path.basename(fname))
252             if os.path.exists(pathname):
253                 return pathname
254
255     # If not found, just return the standard, unchanged path
256     return fname
257
258 def ReadFile(fname, binary=True):
259     """Read and return the contents of a file.
260
261     Args:
262       fname: path to filename to read, where ## signifiies the chroot.
263
264     Returns:
265       data read from file, as a string.
266     """
267     with open(Filename(fname), binary and 'rb' or 'r') as fd:
268         data = fd.read()
269     #self._out.Info("Read file '%s' size %d (%#0x)" %
270                    #(fname, len(data), len(data)))
271     return data
272
273 def WriteFile(fname, data):
274     """Write data into a file.
275
276     Args:
277         fname: path to filename to write
278         data: data to write to file, as a string
279     """
280     #self._out.Info("Write file '%s' size %d (%#0x)" %
281                    #(fname, len(data), len(data)))
282     with open(Filename(fname), 'wb') as fd:
283         fd.write(data)
284
285 def GetBytes(byte, size):
286     """Get a string of bytes of a given size
287
288     This handles the unfortunate different between Python 2 and Python 2.
289
290     Args:
291         byte: Numeric byte value to use
292         size: Size of bytes/string to return
293
294     Returns:
295         A bytes type with 'byte' repeated 'size' times
296     """
297     if sys.version_info[0] >= 3:
298         data = bytes([byte]) * size
299     else:
300         data = chr(byte) * size
301     return data
302
303 def ToUnicode(val):
304     """Make sure a value is a unicode string
305
306     This allows some amount of compatibility between Python 2 and Python3. For
307     the former, it returns a unicode object.
308
309     Args:
310         val: string or unicode object
311
312     Returns:
313         unicode version of val
314     """
315     if sys.version_info[0] >= 3:
316         return val
317     return val if isinstance(val, unicode) else val.decode('utf-8')
318
319 def FromUnicode(val):
320     """Make sure a value is a non-unicode string
321
322     This allows some amount of compatibility between Python 2 and Python3. For
323     the former, it converts a unicode object to a string.
324
325     Args:
326         val: string or unicode object
327
328     Returns:
329         non-unicode version of val
330     """
331     if sys.version_info[0] >= 3:
332         return val
333     return val if isinstance(val, str) else val.encode('utf-8')
334
335 def ToByte(ch):
336     """Convert a character to an ASCII value
337
338     This is useful because in Python 2 bytes is an alias for str, but in
339     Python 3 they are separate types. This function converts the argument to
340     an ASCII value in either case.
341
342     Args:
343         ch: A string (Python 2) or byte (Python 3) value
344
345     Returns:
346         integer ASCII value for ch
347     """
348     return ord(ch) if type(ch) == str else ch
349
350 def ToChar(byte):
351     """Convert a byte to a character
352
353     This is useful because in Python 2 bytes is an alias for str, but in
354     Python 3 they are separate types. This function converts an ASCII value to
355     a value with the appropriate type in either case.
356
357     Args:
358         byte: A byte or str value
359     """
360     return chr(byte) if type(byte) != str else byte
361
362 def ToChars(byte_list):
363     """Convert a list of bytes to a str/bytes type
364
365     Args:
366         byte_list: List of ASCII values representing the string
367
368     Returns:
369         string made by concatenating all the ASCII values
370     """
371     return ''.join([chr(byte) for byte in byte_list])
372
373 def ToBytes(string):
374     """Convert a str type into a bytes type
375
376     Args:
377         string: string to convert
378
379     Returns:
380         Python 3: A bytes type
381         Python 2: A string type
382     """
383     if sys.version_info[0] >= 3:
384         return string.encode('utf-8')
385     return string
386
387 def ToString(bval):
388     """Convert a bytes type into a str type
389
390     Args:
391         bval: bytes value to convert
392
393     Returns:
394         Python 3: A bytes type
395         Python 2: A string type
396     """
397     return bval.decode('utf-8')
398
399 def Compress(indata, algo, with_header=True):
400     """Compress some data using a given algorithm
401
402     Note that for lzma this uses an old version of the algorithm, not that
403     provided by xz.
404
405     This requires 'lz4' and 'lzma_alone' tools. It also requires an output
406     directory to be previously set up, by calling PrepareOutputDir().
407
408     Args:
409         indata: Input data to compress
410         algo: Algorithm to use ('none', 'gzip', 'lz4' or 'lzma')
411
412     Returns:
413         Compressed data
414     """
415     if algo == 'none':
416         return indata
417     fname = GetOutputFilename('%s.comp.tmp' % algo)
418     WriteFile(fname, indata)
419     if algo == 'lz4':
420         data = Run('lz4', '--no-frame-crc', '-c', fname, binary=True)
421     # cbfstool uses a very old version of lzma
422     elif algo == 'lzma':
423         outfname = GetOutputFilename('%s.comp.otmp' % algo)
424         Run('lzma_alone', 'e', fname, outfname, '-lc1', '-lp0', '-pb0', '-d8')
425         data = ReadFile(outfname)
426     elif algo == 'gzip':
427         data = Run('gzip', '-c', fname, binary=True)
428     else:
429         raise ValueError("Unknown algorithm '%s'" % algo)
430     if with_header:
431         hdr = struct.pack('<I', len(data))
432         data = hdr + data
433     return data
434
435 def Decompress(indata, algo, with_header=True):
436     """Decompress some data using a given algorithm
437
438     Note that for lzma this uses an old version of the algorithm, not that
439     provided by xz.
440
441     This requires 'lz4' and 'lzma_alone' tools. It also requires an output
442     directory to be previously set up, by calling PrepareOutputDir().
443
444     Args:
445         indata: Input data to decompress
446         algo: Algorithm to use ('none', 'gzip', 'lz4' or 'lzma')
447
448     Returns:
449         Compressed data
450     """
451     if algo == 'none':
452         return indata
453     if with_header:
454         data_len = struct.unpack('<I', indata[:4])[0]
455         indata = indata[4:4 + data_len]
456     fname = GetOutputFilename('%s.decomp.tmp' % algo)
457     with open(fname, 'wb') as fd:
458         fd.write(indata)
459     if algo == 'lz4':
460         data = Run('lz4', '-dc', fname, binary=True)
461     elif algo == 'lzma':
462         outfname = GetOutputFilename('%s.decomp.otmp' % algo)
463         Run('lzma_alone', 'd', fname, outfname)
464         data = ReadFile(outfname, binary=True)
465     elif algo == 'gzip':
466         data = Run('gzip', '-cd', fname, binary=True)
467     else:
468         raise ValueError("Unknown algorithm '%s'" % algo)
469     return data
470
471 CMD_CREATE, CMD_DELETE, CMD_ADD, CMD_REPLACE, CMD_EXTRACT = range(5)
472
473 IFWITOOL_CMDS = {
474     CMD_CREATE: 'create',
475     CMD_DELETE: 'delete',
476     CMD_ADD: 'add',
477     CMD_REPLACE: 'replace',
478     CMD_EXTRACT: 'extract',
479     }
480
481 def RunIfwiTool(ifwi_file, cmd, fname=None, subpart=None, entry_name=None):
482     """Run ifwitool with the given arguments:
483
484     Args:
485         ifwi_file: IFWI file to operation on
486         cmd: Command to execute (CMD_...)
487         fname: Filename of file to add/replace/extract/create (None for
488             CMD_DELETE)
489         subpart: Name of sub-partition to operation on (None for CMD_CREATE)
490         entry_name: Name of directory entry to operate on, or None if none
491     """
492     args = ['ifwitool', ifwi_file]
493     args.append(IFWITOOL_CMDS[cmd])
494     if fname:
495         args += ['-f', fname]
496     if subpart:
497         args += ['-n', subpart]
498     if entry_name:
499         args += ['-d', '-e', entry_name]
500     Run(*args)
501
502 def ToHex(val):
503     """Convert an integer value (or None) to a string
504
505     Returns:
506         hex value, or 'None' if the value is None
507     """
508     return 'None' if val is None else '%#x' % val
509
510 def ToHexSize(val):
511     """Return the size of an object in hex
512
513     Returns:
514         hex value of size, or 'None' if the value is None
515     """
516     return 'None' if val is None else '%#x' % len(val)