• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Lint as: python2, python3
2# Copyright 2017 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""
7Convenience functions for use by tests or whomever.
8"""
9
10# pylint: disable=missing-docstring
11
12from __future__ import absolute_import
13from __future__ import division
14from __future__ import print_function
15
16import base64
17import collections
18import errno
19import glob
20import json
21import logging
22import math
23import multiprocessing
24import os
25import platform
26import re
27import shutil
28import signal
29import string
30import subprocess
31import tempfile
32import time
33import uuid
34
35from autotest_lib.client.common_lib import error
36from autotest_lib.client.common_lib import magic
37from autotest_lib.client.common_lib import utils
38from autotest_lib.client.common_lib.cros import cros_config
39
40from autotest_lib.client.common_lib.utils import *
41import six
42from six.moves import map
43from six.moves import range
44from six.moves import zip
45
46
47def grep(pattern, file):
48    """
49    This is mainly to fix the return code inversion from grep
50    Also handles compressed files.
51
52    returns 1 if the pattern is present in the file, 0 if not.
53    """
54    command = 'grep "%s" > /dev/null' % pattern
55    ret = cat_file_to_cmd(file, command, ignore_status=True)
56    return not ret
57
58
59def cat_file_to_cmd(file, command, ignore_status=0, return_output=False):
60    """
61    equivalent to 'cat file | command' but knows to use
62    zcat or bzcat if appropriate
63    """
64    if not os.path.isfile(file):
65        raise NameError('invalid file %s to cat to command %s'
66                % (file, command))
67
68    if return_output:
69        run_cmd = utils.system_output
70    else:
71        run_cmd = utils.system
72
73    if magic.guess_type(file) == 'application/x-bzip2':
74        cat = 'bzcat'
75    elif magic.guess_type(file) == 'application/x-gzip':
76        cat = 'zcat'
77    else:
78        cat = 'cat'
79    return run_cmd('%s %s | %s' % (cat, file, command),
80                   ignore_status=ignore_status)
81
82
83def extract_tarball_to_dir(tarball, dir):
84    """
85    Extract a tarball to a specified directory name instead of whatever
86    the top level of a tarball is - useful for versioned directory names, etc
87    """
88    if os.path.exists(dir):
89        if os.path.isdir(dir):
90            shutil.rmtree(dir)
91        else:
92            os.remove(dir)
93    pwd = os.getcwd()
94    os.chdir(os.path.dirname(os.path.abspath(dir)))
95    newdir = extract_tarball(tarball)
96    os.rename(newdir, dir)
97    os.chdir(pwd)
98
99
100def extract_tarball(tarball):
101    """Returns the directory extracted by the tarball."""
102    extracted = cat_file_to_cmd(tarball, 'tar xvf - 2>/dev/null',
103                                    return_output=True).splitlines()
104
105    dir = None
106
107    for line in extracted:
108        if line.startswith('./'):
109            line = line[2:]
110        if not line or line == '.':
111            continue
112        topdir = line.split('/')[0]
113        if os.path.isdir(topdir):
114            if dir:
115                assert(dir == topdir), 'tarball must be a a single directory'
116            else:
117                dir = topdir
118    if dir:
119        return dir
120    else:
121        raise NameError('extracting tarball produced no dir')
122
123
124def force_copy(src, dest):
125    """Replace dest with a new copy of src, even if it exists"""
126    if os.path.isfile(dest):
127        os.remove(dest)
128    if os.path.isdir(dest):
129        dest = os.path.join(dest, os.path.basename(src))
130    shutil.copyfile(src, dest)
131    return dest
132
133
134def file_contains_pattern(file, pattern):
135    """Return true if file contains the specified egrep pattern"""
136    if not os.path.isfile(file):
137        raise NameError('file %s does not exist' % file)
138    return not utils.system('egrep -q "' + pattern + '" ' + file,
139                            ignore_status=True)
140
141
142def list_grep(list, pattern):
143    """True if any item in list matches the specified pattern."""
144    compiled = re.compile(pattern)
145    for line in list:
146        match = compiled.search(line)
147        if (match):
148            return 1
149    return 0
150
151
152def get_os_vendor():
153    """Try to guess what's the os vendor
154    """
155    if os.path.isfile('/etc/SuSE-release'):
156        return 'SUSE'
157
158    issue = '/etc/issue'
159
160    if not os.path.isfile(issue):
161        return 'Unknown'
162
163    if file_contains_pattern(issue, 'Red Hat'):
164        return 'Red Hat'
165    elif file_contains_pattern(issue, 'Fedora'):
166        return 'Fedora Core'
167    elif file_contains_pattern(issue, 'SUSE'):
168        return 'SUSE'
169    elif file_contains_pattern(issue, 'Ubuntu'):
170        return 'Ubuntu'
171    elif file_contains_pattern(issue, 'Debian'):
172        return 'Debian'
173    else:
174        return 'Unknown'
175
176
177def get_cc():
178    try:
179        return os.environ['CC']
180    except KeyError:
181        return 'gcc'
182
183
184def get_vmlinux():
185    """Return the full path to vmlinux
186
187    Ahem. This is crap. Pray harder. Bad Martin.
188    """
189    vmlinux = '/boot/vmlinux-%s' % utils.system_output('uname -r')
190    if os.path.isfile(vmlinux):
191        return vmlinux
192    vmlinux = '/lib/modules/%s/build/vmlinux' % utils.system_output('uname -r')
193    if os.path.isfile(vmlinux):
194        return vmlinux
195    return None
196
197
198def get_systemmap():
199    """Return the full path to System.map
200
201    Ahem. This is crap. Pray harder. Bad Martin.
202    """
203    map = '/boot/System.map-%s' % utils.system_output('uname -r')
204    if os.path.isfile(map):
205        return map
206    map = '/lib/modules/%s/build/System.map' % utils.system_output('uname -r')
207    if os.path.isfile(map):
208        return map
209    return None
210
211
212def get_modules_dir():
213    """Return the modules dir for the running kernel version"""
214    kernel_version = utils.system_output('uname -r')
215    return '/lib/modules/%s/kernel' % kernel_version
216
217
218_CPUINFO_RE = re.compile(r'^(?P<key>[^\t]*)\t*: ?(?P<value>.*)$')
219
220
221def get_cpuinfo():
222    """Read /proc/cpuinfo and convert to a list of dicts."""
223    cpuinfo = []
224    with open('/proc/cpuinfo', 'r') as f:
225        cpu = {}
226        for line in f:
227            line = line.strip()
228            if not line:
229                cpuinfo.append(cpu)
230                cpu = {}
231                continue
232            match = _CPUINFO_RE.match(line)
233            cpu[match.group('key')] = match.group('value')
234        if cpu:
235            # cpuinfo usually ends in a blank line, so this shouldn't happen.
236            cpuinfo.append(cpu)
237    return cpuinfo
238
239
240def get_cpu_arch():
241    """Work out which CPU architecture we're running on"""
242
243    # Using 'uname -m' should be a very portable way to do this since the
244    # format is pretty standard.
245    machine_name = utils.system_output('uname -m').strip()
246
247    # Apparently ARM64 and ARM have both historically returned the string 'arm'
248    # here so continue the tradition.  Use startswith() because:
249    # - On most of our arm devices we'll actually see the string armv7l.
250    # - In theory the machine name could include a suffix for endianness.
251    if machine_name.startswith('aarch64') or machine_name.startswith('arm'):
252        return 'arm'
253
254    # Historically we _have_ treated x86_64 and i386 separately.
255    if machine_name in ('x86_64', 'i386'):
256        return machine_name
257
258    raise error.TestError('unsupported machine type %s' % machine_name)
259
260
261def get_arm_soc_family_from_devicetree():
262    """
263    Work out which ARM SoC we're running on based on the 'compatible' property
264    of the base node of devicetree, if it exists.
265    """
266    devicetree_compatible = '/sys/firmware/devicetree/base/compatible'
267    if not os.path.isfile(devicetree_compatible):
268        return None
269    f = open(devicetree_compatible, 'r')
270    compatible = f.read().split(chr(0))
271    f.close()
272    if list_grep(compatible, '^rockchip,'):
273        return 'rockchip'
274    elif list_grep(compatible, '^mediatek,'):
275        return 'mediatek'
276    elif list_grep(compatible, '^qcom,'):
277        return 'qualcomm'
278    return None
279
280
281def get_arm_soc_family():
282    """Work out which ARM SoC we're running on"""
283    family = get_arm_soc_family_from_devicetree()
284    if family is not None:
285        return family
286
287    f = open('/proc/cpuinfo', 'r')
288    cpuinfo = f.readlines()
289    f.close()
290    if list_grep(cpuinfo, 'EXYNOS5'):
291        return 'exynos5'
292    elif list_grep(cpuinfo, 'Tegra'):
293        return 'tegra'
294    elif list_grep(cpuinfo, 'Rockchip'):
295        return 'rockchip'
296    return 'arm'
297
298
299def get_cpu_soc_family():
300    """Like get_cpu_arch, but for ARM, returns the SoC family name"""
301    f = open('/proc/cpuinfo', 'r')
302    cpuinfo = f.readlines()
303    f.close()
304    family = get_cpu_arch()
305    if family == 'arm':
306        family = get_arm_soc_family()
307    if list_grep(cpuinfo, '^vendor_id.*:.*AMD'):
308        family = 'amd'
309    return family
310
311
312INTEL_UARCH_TABLE = {
313    '06_4C': 'Airmont',
314    '06_1C': 'Atom',
315    '06_26': 'Atom',
316    '06_27': 'Atom',
317    '06_35': 'Atom',
318    '06_36': 'Atom',
319    '06_3D': 'Broadwell',
320    '06_47': 'Broadwell',
321    '06_4F': 'Broadwell',
322    '06_56': 'Broadwell',
323    '06_A5': 'Comet Lake',
324    '06_A6': 'Comet Lake',
325    '06_0D': 'Dothan',
326    '06_5C': 'Goldmont',
327    '06_7A': 'Goldmont',
328    '06_3C': 'Haswell',
329    '06_45': 'Haswell',
330    '06_46': 'Haswell',
331    '06_3F': 'Haswell-E',
332    '06_7D': 'Ice Lake',
333    '06_7E': 'Ice Lake',
334    '06_3A': 'Ivy Bridge',
335    '06_3E': 'Ivy Bridge-E',
336    '06_8E': 'Kaby Lake',
337    '06_9E': 'Kaby Lake',
338    '06_0F': 'Merom',
339    '06_16': 'Merom',
340    '06_17': 'Nehalem',
341    '06_1A': 'Nehalem',
342    '06_1D': 'Nehalem',
343    '06_1E': 'Nehalem',
344    '06_1F': 'Nehalem',
345    '06_2E': 'Nehalem',
346    '0F_03': 'Prescott',
347    '0F_04': 'Prescott',
348    '0F_06': 'Presler',
349    '06_2A': 'Sandy Bridge',
350    '06_2D': 'Sandy Bridge',
351    '06_37': 'Silvermont',
352    '06_4A': 'Silvermont',
353    '06_4D': 'Silvermont',
354    '06_5A': 'Silvermont',
355    '06_5D': 'Silvermont',
356    '06_4E': 'Skylake',
357    '06_5E': 'Skylake',
358    '06_55': 'Skylake',
359    '06_8C': 'Tiger Lake',
360    '06_8D': 'Tiger Lake',
361    '06_86': 'Tremont',
362    '06_96': 'Tremont',
363    '06_9C': 'Tremont',
364    '06_25': 'Westmere',
365    '06_2C': 'Westmere',
366    '06_2F': 'Westmere',
367}
368
369
370def get_intel_cpu_uarch(numeric=False):
371    """Return the Intel microarchitecture we're running on, or None.
372
373    Returns None if this is not an Intel CPU. Returns the family and model as
374    underscore-separated hex (per Intel manual convention) if the uarch is not
375    known, or if numeric is True.
376    """
377    if not get_current_kernel_arch().startswith('x86'):
378        return None
379    cpuinfo = get_cpuinfo()[0]
380    if cpuinfo['vendor_id'] != 'GenuineIntel':
381        return None
382    family_model = '%02X_%02X' % (int(cpuinfo['cpu family']),
383                                  int(cpuinfo['model']))
384    if numeric:
385        return family_model
386    return INTEL_UARCH_TABLE.get(family_model, family_model)
387
388
389INTEL_SILVERMONT_BCLK_TABLE = [83333, 100000, 133333, 116667, 80000];
390
391
392def get_intel_bclk_khz():
393    """Return Intel CPU base clock.
394
395    This only worked with SandyBridge (released in 2011) or newer. Older CPU has
396    133 MHz bclk. See turbostat code for implementation that also works with
397    older CPU. https://git.io/vpyKT
398    """
399    if get_intel_cpu_uarch() == 'Silvermont':
400        MSR_FSB_FREQ = 0xcd
401        return INTEL_SILVERMONT_BCLK_TABLE[utils.rdmsr(MSR_FSB_FREQ) & 0xf]
402    return 100000
403
404
405def get_current_kernel_arch():
406    """Get the machine architecture, now just a wrap of 'uname -m'."""
407    return os.popen('uname -m').read().rstrip()
408
409
410def count_cpus():
411    """number of CPUs in the local machine according to /proc/cpuinfo"""
412    try:
413       return multiprocessing.cpu_count()
414    except Exception:
415       logging.exception('can not get cpu count from'
416                        ' multiprocessing.cpu_count()')
417    cpuinfo = get_cpuinfo()
418    # Returns at least one cpu. Check comment #1 in crosbug.com/p/9582.
419    return len(cpuinfo) or 1
420
421
422def cpu_online_map():
423    """
424    Check out the available cpu online map
425    """
426    cpuinfo = get_cpuinfo()
427    cpus = []
428    for cpu in cpuinfo:
429        cpus.append(cpu['processor'])  # grab cpu number
430    return cpus
431
432
433# Returns total memory in kb
434def read_from_meminfo(key):
435    meminfo = utils.system_output('grep %s /proc/meminfo' % key)
436    return int(re.search(r'\d+', meminfo).group(0))
437
438
439def memtotal():
440    return read_from_meminfo('MemTotal')
441
442
443def freememtotal():
444    return read_from_meminfo('MemFree')
445
446def usable_memtotal():
447    # Reserved 5% for OS use
448    return int(read_from_meminfo('MemFree') * 0.95)
449
450def swaptotal():
451    return read_from_meminfo('SwapTotal')
452
453def rounded_memtotal():
454    # Get total of all physical mem, in kbytes
455    usable_kbytes = memtotal()
456    # usable_kbytes is system's usable DRAM in kbytes,
457    #   as reported by memtotal() from device /proc/meminfo memtotal
458    #   after Linux deducts 1.5% to 5.1% for system table overhead
459    # Undo the unknown actual deduction by rounding up
460    #   to next small multiple of a big power-of-two
461    #   eg  12GB - 5.1% gets rounded back up to 12GB
462    mindeduct = 0.015  # 1.5 percent
463    maxdeduct = 0.055  # 5.5 percent
464    # deduction range 1.5% .. 5.5% supports physical mem sizes
465    #    6GB .. 12GB in steps of .5GB
466    #   12GB .. 24GB in steps of 1 GB
467    #   24GB .. 48GB in steps of 2 GB ...
468    # Finer granularity in physical mem sizes would require
469    #   tighter spread between min and max possible deductions
470
471    # increase mem size by at least min deduction, without rounding
472    min_kbytes = int(usable_kbytes / (1.0 - mindeduct))
473    # increase mem size further by 2**n rounding, by 0..roundKb or more
474    round_kbytes = int(usable_kbytes / (1.0 - maxdeduct)) - min_kbytes
475    # find least binary roundup 2**n that covers worst-cast roundKb
476    mod2n = 1 << int(math.ceil(math.log(round_kbytes, 2)))
477    # have round_kbytes <= mod2n < round_kbytes*2
478    # round min_kbytes up to next multiple of mod2n
479    phys_kbytes = min_kbytes + mod2n - 1
480    phys_kbytes = phys_kbytes - (phys_kbytes % mod2n)  # clear low bits
481    return phys_kbytes
482
483
484_MEMINFO_RE = re.compile('^(\w+)(\(\w+\))?:\s+(\d+)')
485
486
487def get_meminfo():
488    """Returns a namedtuple of pairs from /proc/meminfo.
489
490    Example /proc/meminfo snippets:
491        MemTotal:        2048000 kB
492        Active(anon):     409600 kB
493    Example usage:
494        meminfo = utils.get_meminfo()
495        print meminfo.Active_anon
496    """
497    info = {}
498    with _open_file('/proc/meminfo') as f:
499        for line in f:
500            m = _MEMINFO_RE.match(line)
501            if m:
502                if m.group(2):
503                    name = m.group(1) + '_' + m.group(2)[1:-1]
504                else:
505                    name = m.group(1)
506                info[name] = int(m.group(3))
507    return collections.namedtuple('MemInfo', list(info.keys()))(**info)
508
509
510def sysctl(key, value=None):
511    """Generic implementation of sysctl, to read and write.
512
513    @param key: A location under /proc/sys
514    @param value: If not None, a value to write into the sysctl.
515
516    @return The single-line sysctl value as a string.
517    """
518    path = '/proc/sys/%s' % key
519    if value is not None:
520        utils.write_one_line(path, str(value))
521    return utils.read_one_line(path)
522
523
524def sysctl_kernel(key, value=None):
525    """(Very) partial implementation of sysctl, for kernel params"""
526    if value is not None:
527        # write
528        utils.write_one_line('/proc/sys/kernel/%s' % key, str(value))
529    else:
530        # read
531        out = utils.read_one_line('/proc/sys/kernel/%s' % key)
532        return int(re.search(r'\d+', out).group(0))
533
534
535def get_num_allocated_file_handles():
536    """
537    Returns the number of currently allocated file handles.
538
539    Gets this information by parsing /proc/sys/fs/file-nr.
540    See https://www.kernel.org/doc/Documentation/sysctl/fs.txt
541    for details on this file.
542    """
543    with _open_file('/proc/sys/fs/file-nr') as f:
544        line = f.readline()
545    allocated_handles = int(line.split()[0])
546    return allocated_handles
547
548
549def dump_object(object):
550    """Dump an object's attributes and methods
551
552    kind of like dir()
553    """
554    for item in six.iteritems(object.__dict__):
555        print(item)
556        try:
557            (key, value) = item
558            dump_object(value)
559        except:
560            continue
561
562
563def environ(env_key):
564    """return the requested environment variable, or '' if unset"""
565    if (env_key in os.environ):
566        return os.environ[env_key]
567    else:
568        return ''
569
570
571def prepend_path(newpath, oldpath):
572    """prepend newpath to oldpath"""
573    if (oldpath):
574        return newpath + ':' + oldpath
575    else:
576        return newpath
577
578
579def append_path(oldpath, newpath):
580    """append newpath to oldpath"""
581    if (oldpath):
582        return oldpath + ':' + newpath
583    else:
584        return newpath
585
586
587_TIME_OUTPUT_RE = re.compile(
588        r'([\d\.]*)user ([\d\.]*)system '
589        r'(\d*):([\d\.]*)elapsed (\d*)%CPU')
590
591
592def to_seconds(time_string):
593    """Converts a string in M+:SS.SS format to S+.SS"""
594    elts = time_string.split(':')
595    if len(elts) == 1:
596        return time_string
597    return str(int(elts[0]) * 60 + float(elts[1]))
598
599
600_TIME_OUTPUT_RE_2 = re.compile(r'(.*?)user (.*?)system (.*?)elapsed')
601
602
603def running_config():
604    """
605    Return path of config file of the currently running kernel
606    """
607    version = utils.system_output('uname -r')
608    for config in ('/proc/config.gz', \
609                   '/boot/config-%s' % version,
610                   '/lib/modules/%s/build/.config' % version):
611        if os.path.isfile(config):
612            return config
613    return None
614
615
616def check_for_kernel_feature(feature):
617    config = running_config()
618
619    if not config:
620        raise TypeError("Can't find kernel config file")
621
622    if magic.guess_type(config) == 'application/x-gzip':
623        grep = 'zgrep'
624    else:
625        grep = 'grep'
626    grep += ' ^CONFIG_%s= %s' % (feature, config)
627
628    if not utils.system_output(grep, ignore_status=True):
629        raise ValueError("Kernel doesn't have a %s feature" % (feature))
630
631
632def check_glibc_ver(ver):
633    try:
634        glibc_ver = subprocess.check_output("ldd --version", shell=True)
635    except subprocess.CalledProcessError:
636        # To mimic previous behavior, if the command errors set the result to
637        # an empty str
638        glibc_ver = ''
639    glibc_ver = glibc_ver.splitlines()[0].decode()
640    glibc_ver = re.search(r'(\d+\.\d+(\.\d+)?)', glibc_ver).group()
641    if utils.compare_versions(glibc_ver, ver) == -1:
642        raise error.TestError("Glibc too old (%s). Glibc >= %s is needed." %
643                              (glibc_ver, ver))
644
645def check_kernel_ver(ver):
646    kernel_ver = utils.system_output('uname -r')
647    kv_tmp = re.split(r'[-]', kernel_ver)[0:3]
648    # In compare_versions, if v1 < v2, return value == -1
649    if utils.compare_versions(kv_tmp[0], ver) == -1:
650        raise error.TestError("Kernel too old (%s). Kernel > %s is needed." %
651                              (kernel_ver, ver))
652
653
654def numa_nodes():
655    node_paths = glob.glob('/sys/devices/system/node/node*')
656    nodes = [int(re.sub(r'.*node(\d+)', r'\1', x)) for x in node_paths]
657    return (sorted(nodes))
658
659
660# Return the kernel version and build timestamp.
661def running_os_release():
662    return os.uname()[2:4]
663
664
665def running_os_ident():
666    (version, timestamp) = running_os_release()
667    return version + '::' + timestamp
668
669
670def freespace(path):
671    """Return the disk free space, in bytes"""
672    s = os.statvfs(path)
673    return s.f_bavail * s.f_bsize
674
675
676_DISK_PARTITION_3_RE = re.compile(r'^(/dev/hd[a-z]+)3', re.M)
677
678
679def get_disk_size(disk_name):
680    """
681    Return size of disk in byte. Return 0 in Error Case
682
683    @param disk_name: disk name to find size
684    """
685    device = os.path.basename(disk_name)
686    with open('/proc/partitions') as f:
687        lines = f.readlines()
688    for line in lines:
689        try:
690            _, _, blocks, name = re.split(r' +', line.strip())
691        except ValueError:
692            continue
693        if name == device:
694            return 1024 * int(blocks)
695    return 0
696
697
698def get_disk_size_gb(disk_name):
699    """
700    Return size of disk in GB (10^9). Return 0 in Error Case
701
702    @param disk_name: disk name to find size
703    """
704    return int(get_disk_size(disk_name) / (10.0 ** 9) + 0.5)
705
706
707def get_disk_model(disk_name):
708    """
709    Return model name for internal storage device
710
711    @param disk_name: disk name to find model
712    """
713    cmd1 = 'udevadm info --query=property --name=%s' % disk_name
714    cmd2 = 'grep -E "ID_(NAME|MODEL)="'
715    cmd3 = 'cut -f 2 -d"="'
716    cmd = ' | '.join([cmd1, cmd2, cmd3])
717    return utils.system_output(cmd)
718
719
720_DISK_DEV_RE = re.compile(r'/dev/sd[a-z]|'
721                          r'/dev/mmcblk[0-9]+|'
722                          r'/dev/nvme[0-9]+n[0-9]+')
723
724
725def get_disk_from_filename(filename):
726    """
727    Return the disk device the filename is on.
728    If the file is on tmpfs or other special file systems,
729    return None.
730
731    @param filename: name of file, full path.
732    """
733
734    if not os.path.exists(filename):
735        raise error.TestError('file %s missing' % filename)
736
737    if filename[0] != '/':
738        raise error.TestError('This code works only with full path')
739
740    m = _DISK_DEV_RE.match(filename)
741    while not m:
742        if filename[0] != '/':
743            return None
744        if filename == '/dev/root':
745            cmd = 'rootdev -d -s'
746        elif filename.startswith('/dev/mapper'):
747            cmd = 'dmsetup table "%s"' % os.path.basename(filename)
748            dmsetup_output = utils.system_output(cmd).split(' ')
749            if dmsetup_output[2] == 'verity':
750                maj_min = dmsetup_output[4]
751            elif dmsetup_output[2] == 'crypt':
752                maj_min = dmsetup_output[6]
753            cmd = 'realpath "/dev/block/%s"' % maj_min
754        elif filename.startswith('/dev/loop'):
755            cmd = 'losetup -O BACK-FILE "%s" | tail -1' % filename
756        else:
757            cmd = 'df "%s" | tail -1 | cut -f 1 -d" "' % filename
758        filename = utils.system_output(cmd)
759        m = _DISK_DEV_RE.match(filename)
760    return m.group(0)
761
762
763def get_disk_firmware_version(disk_name):
764    """
765    Return firmware version for internal storage device. (empty string for eMMC)
766
767    @param disk_name: disk name to find model
768    """
769    cmd1 = 'udevadm info --query=property --name=%s' % disk_name
770    cmd2 = 'grep -E "ID_REVISION="'
771    cmd3 = 'cut -f 2 -d"="'
772    cmd = ' | '.join([cmd1, cmd2, cmd3])
773    return utils.system_output(cmd)
774
775
776def is_disk_nvme(disk_name):
777    """
778    Return true if disk is a nvme device, return false otherwise
779
780    @param disk_name: disk name to check
781    """
782    return re.match('/dev/nvme[0-9]+n[0-9]+', disk_name)
783
784
785def is_disk_scsi(disk_name):
786    """
787    Return true if disk is a scsi device, return false otherwise
788
789    @param disk_name: disk name check
790    """
791    return re.match('/dev/sd[a-z]+', disk_name)
792
793
794def is_disk_harddisk(disk_name):
795    """
796    Return true if disk is a harddisk, return false otherwise
797
798    @param disk_name: disk name check
799    """
800    cmd1 = 'udevadm info --query=property --name=%s' % disk_name
801    cmd2 = 'grep -E "ID_ATA_ROTATION_RATE_RPM="'
802    cmd3 = 'cut -f 2 -d"="'
803    cmd = ' | '.join([cmd1, cmd2, cmd3])
804
805    rtt = utils.system_output(cmd)
806
807    # eMMC will not have this field; rtt == ''
808    # SSD will have zero rotation rate; rtt == '0'
809    # For harddisk rtt > 0
810    return rtt and int(rtt) > 0
811
812def concat_partition(disk_name, partition_number):
813    """
814    Return the name of a partition:
815    sda, 3 --> sda3
816    mmcblk0, 3 --> mmcblk0p3
817
818    @param disk_name: diskname string
819    @param partition_number: integer
820    """
821    if disk_name.endswith(tuple(str(i) for i in range(0, 10))):
822        sep = 'p'
823    else:
824        sep = ''
825    return disk_name + sep + str(partition_number)
826
827def verify_hdparm_feature(disk_name, feature):
828    """
829    Check for feature support for SCSI disk using hdparm
830
831    @param disk_name: target disk
832    @param feature: hdparm output string of the feature
833    """
834    cmd = 'hdparm -I %s | grep -q "%s"' % (disk_name, feature)
835    ret = utils.system(cmd, ignore_status=True)
836    if ret == 0:
837        return True
838    elif ret == 1:
839        return False
840    else:
841        raise error.TestFail('Error running command %s' % cmd)
842
843def get_nvme_id_ns_feature(disk_name, feature):
844    """
845    Return feature value for NVMe disk using nvme id-ns
846
847    @param disk_name: target disk
848    @param feature: output string of the feature
849    """
850    cmd = "nvme id-ns -n 1 %s | grep %s" % (disk_name, feature)
851    feat = utils.system_output(cmd, ignore_status=True)
852    if not feat:
853        return 'None'
854    start = feat.find(':')
855    value = feat[start+2:]
856    return value
857
858def get_storage_error_msg(disk_name, reason):
859    """
860    Get Error message for storage test which include disk model.
861    and also include the firmware version for the SCSI disk
862
863    @param disk_name: target disk
864    @param reason: Reason of the error.
865    """
866
867    msg = reason
868
869    model = get_disk_model(disk_name)
870    msg += ' Disk model: %s' % model
871
872    if is_disk_scsi(disk_name):
873        fw = get_disk_firmware_version(disk_name)
874        msg += ' firmware: %s' % fw
875
876    return msg
877
878
879_IOSTAT_FIELDS = ('transfers_per_s', 'read_kb_per_s', 'written_kb_per_s',
880                  'read_kb', 'written_kb')
881_IOSTAT_RE = re.compile('ALL' + len(_IOSTAT_FIELDS) * r'\s+([\d\.]+)')
882
883def get_storage_statistics(device=None):
884    """
885    Fetches statistics for a storage device.
886
887    Using iostat(1) it retrieves statistics for a device since last boot.  See
888    the man page for iostat(1) for details on the different fields.
889
890    @param device: Path to a block device. Defaults to the device where root
891            is mounted.
892
893    @returns a dict mapping each field to its statistic.
894
895    @raises ValueError: If the output from iostat(1) can not be parsed.
896    """
897    if device is None:
898        device = get_root_device()
899    cmd = 'iostat -d -k -g ALL -H %s' % device
900    output = utils.system_output(cmd, ignore_status=True)
901    match = _IOSTAT_RE.search(output)
902    if not match:
903        raise ValueError('Unable to get iostat for %s' % device)
904    return dict(list(zip(_IOSTAT_FIELDS, list(map(float, match.groups())))))
905
906
907def load_module(module_name, params=None):
908    # Checks if a module has already been loaded
909    if module_is_loaded(module_name):
910        return False
911
912    cmd = '/sbin/modprobe ' + module_name
913    if params:
914        cmd += ' ' + params
915    utils.system(cmd)
916    return True
917
918
919def unload_module(module_name):
920    """
921    Removes a module. Handles dependencies. If even then it's not possible
922    to remove one of the modules, it will trhow an error.CmdError exception.
923
924    @param module_name: Name of the module we want to remove.
925    """
926    module_name = module_name.replace('-', '_')
927    l_raw = utils.system_output("/bin/lsmod").splitlines()
928    lsmod = [x for x in l_raw if x.split()[0] == module_name]
929    if len(lsmod) > 0:
930        line_parts = lsmod[0].split()
931        if len(line_parts) == 4:
932            submodules = line_parts[3].split(",")
933            for submodule in submodules:
934                unload_module(submodule)
935        utils.system("/sbin/modprobe -r %s" % module_name)
936        logging.info("Module %s unloaded", module_name)
937    else:
938        logging.info("Module %s is already unloaded", module_name)
939
940
941def module_is_loaded(module_name):
942    module_name = module_name.replace('-', '_')
943    modules = utils.system_output('/bin/lsmod').splitlines()
944    for module in modules:
945        if module.startswith(module_name) and module[len(module_name)] == ' ':
946            return True
947    return False
948
949
950def ping_default_gateway():
951    """Ping the default gateway."""
952
953    network = open('/etc/sysconfig/network')
954    m = re.search('GATEWAY=(\S+)', network.read())
955
956    if m:
957        gw = m.group(1)
958        cmd = 'ping %s -c 5 > /dev/null' % gw
959        return utils.system(cmd, ignore_status=True)
960
961    raise error.TestError('Unable to find default gateway')
962
963
964def drop_caches():
965    """Writes back all dirty pages to disk and clears all the caches."""
966    utils.system("sync")
967    # We ignore failures here as this will fail on 2.6.11 kernels.
968    utils.system("echo 3 > /proc/sys/vm/drop_caches", ignore_status=True)
969
970
971def set_hwclock(time='system',
972                utc=True,
973                rtc=None,
974                noadjfile=False,
975                ignore_status=False):
976    """Uses the hwclock command to set time of an RTC.
977
978    @param time: Either 'system', meaning use the system time, or a string
979                 to be passed to the --date argument of hwclock.
980    @param utc: Boolean of whether to use UTC or localtime.
981    @param rtc: String to be passed to the --rtc arg of hwclock.
982    @param noadjfile: Boolean of whether to use --noadjfile flag with hwclock.
983    @param ignore_status: Boolean of whether to ignore exit code of hwclock.
984    """
985    cmd = '/sbin/hwclock'
986    if time == 'system':
987        cmd += ' --systohc'
988    else:
989        cmd += ' --set --date "{}"'.format(time)
990    if utc:
991        cmd += ' --utc'
992    else:
993        cmd += ' --localtime'
994    if rtc is not None:
995        cmd += ' --rtc={}'.format(rtc)
996    if noadjfile:
997        cmd += ' --noadjfile'
998    return utils.system(cmd, ignore_status=ignore_status)
999
1000def set_wake_alarm(alarm_time):
1001    """
1002    Set the hardware RTC-based wake alarm to 'alarm_time'.
1003    """
1004    utils.write_one_line('/sys/class/rtc/rtc0/wakealarm', str(alarm_time))
1005
1006
1007_AUTOTEST_CLIENT_PATH = os.path.join(os.path.dirname(__file__), '..')
1008_AMD_PCI_IDS_FILE_PATH = os.path.join(_AUTOTEST_CLIENT_PATH,
1009                                      'bin/amd_pci_ids.json')
1010_INTEL_PCI_IDS_FILE_PATH = os.path.join(_AUTOTEST_CLIENT_PATH,
1011                                        'bin/intel_pci_ids.json')
1012_UI_USE_FLAGS_FILE_PATH = '/etc/ui_use_flags.txt'
1013
1014# Command to check if a package is installed. If the package is not installed
1015# the command shall fail.
1016_CHECK_PACKAGE_INSTALLED_COMMAND =(
1017        "dpkg-query -W -f='${Status}\n' %s | head -n1 | awk '{print $3;}' | "
1018        "grep -q '^installed$'")
1019
1020pciid_to_amd_architecture = {}
1021pciid_to_intel_architecture = {}
1022
1023class Crossystem(object):
1024    """A wrapper for the crossystem utility."""
1025
1026    def __init__(self, client):
1027        self.cros_system_data = {}
1028        self._client = client
1029
1030    def init(self):
1031        self.cros_system_data = {}
1032        (_, fname) = tempfile.mkstemp()
1033        f = open(fname, 'w')
1034        self._client.run('crossystem', stdout_tee=f)
1035        f.close()
1036        text = utils.read_file(fname)
1037        for line in text.splitlines():
1038            assignment_string = line.split('#')[0]
1039            if not assignment_string.count('='):
1040                continue
1041            (name, value) = assignment_string.split('=', 1)
1042            self.cros_system_data[name.strip()] = value.strip()
1043        os.remove(fname)
1044
1045    def __getattr__(self, name):
1046        """
1047        Retrieve a crosssystem attribute.
1048
1049        The call crossystemobject.name() will return the crossystem reported
1050        string.
1051        """
1052        return lambda: self.cros_system_data[name]
1053
1054
1055def get_oldest_pid_by_name(name):
1056    """
1057    Return the oldest pid of a process whose name perfectly matches |name|.
1058
1059    name is an egrep expression, which will be matched against the entire name
1060    of processes on the system.  For example:
1061
1062      get_oldest_pid_by_name('chrome')
1063
1064    on a system running
1065      8600 ?        00:00:04 chrome
1066      8601 ?        00:00:00 chrome
1067      8602 ?        00:00:00 chrome-sandbox
1068
1069    would return 8600, as that's the oldest process that matches.
1070    chrome-sandbox would not be matched.
1071
1072    Arguments:
1073      name: egrep expression to match.  Will be anchored at the beginning and
1074            end of the match string.
1075
1076    Returns:
1077      pid as an integer, or None if one cannot be found.
1078
1079    Raises:
1080      ValueError if pgrep returns something odd.
1081    """
1082    str_pid = utils.system_output('pgrep -o ^%s$' % name,
1083                                  ignore_status=True).rstrip()
1084    if str_pid:
1085        return int(str_pid)
1086
1087
1088def get_oldest_by_name(name):
1089    """Return pid and command line of oldest process whose name matches |name|.
1090
1091    @param name: egrep expression to match desired process name.
1092    @return: A tuple of (pid, command_line) of the oldest process whose name
1093             matches |name|.
1094
1095    """
1096    pid = get_oldest_pid_by_name(name)
1097    if pid:
1098        command_line = utils.system_output('ps -p %i -o command=' % pid,
1099                                           ignore_status=True).rstrip()
1100        return (pid, command_line)
1101
1102
1103def get_chrome_remote_debugging_port():
1104    """Returns remote debugging port for Chrome.
1105
1106    Parse chrome process's command line argument to get the remote debugging
1107    port. if it is 0, look at DevToolsActivePort for the ephemeral port.
1108    """
1109    _, command = get_oldest_by_name('chrome')
1110    matches = re.search('--remote-debugging-port=([0-9]+)', command)
1111    if not matches:
1112      return 0
1113    port = int(matches.group(1))
1114    if port:
1115      return port
1116    with open('/home/chronos/DevToolsActivePort') as f:
1117      return int(f.readline().rstrip())
1118
1119
1120def get_process_list(name, command_line=None):
1121    """
1122    Return the list of pid for matching process |name command_line|.
1123
1124    on a system running
1125      31475 ?    0:06 /opt/google/chrome/chrome --allow-webui-compositing -
1126      31478 ?    0:00 /opt/google/chrome/chrome-sandbox /opt/google/chrome/
1127      31485 ?    0:00 /opt/google/chrome/chrome --type=zygote --log-level=1
1128      31532 ?    1:05 /opt/google/chrome/chrome --type=renderer
1129
1130    get_process_list('chrome')
1131    would return ['31475', '31485', '31532']
1132
1133    get_process_list('chrome', '--type=renderer')
1134    would return ['31532']
1135
1136    Arguments:
1137      name: process name to search for. If command_line is provided, name is
1138            matched against full command line. If command_line is not provided,
1139            name is only matched against the process name.
1140      command line: when command line is passed, the full process command line
1141                    is used for matching.
1142
1143    Returns:
1144      list of PIDs of the matching processes.
1145
1146    """
1147    # TODO(rohitbm) crbug.com/268861
1148    flag = '-x' if not command_line else '-f'
1149    name = '\'%s.*%s\'' % (name, command_line) if command_line else name
1150    str_pid = utils.system_output('pgrep %s %s' % (flag, name),
1151                                  ignore_status=True).rstrip()
1152    return str_pid.split()
1153
1154
1155def nuke_process_by_name(name, with_prejudice=False):
1156    """Tell the oldest process specified by name to exit.
1157
1158    Arguments:
1159      name: process name specifier, as understood by pgrep.
1160      with_prejudice: if True, don't allow for graceful exit.
1161
1162    Raises:
1163      error.AutoservPidAlreadyDeadError: no existing process matches name.
1164    """
1165    try:
1166        pid = get_oldest_pid_by_name(name)
1167    except Exception as e:
1168        logging.error(e)
1169        return
1170    if pid is None:
1171        raise error.AutoservPidAlreadyDeadError('No process matching %s.' %
1172                                                name)
1173    if with_prejudice:
1174        utils.nuke_pid(pid, [signal.SIGKILL])
1175    else:
1176        utils.nuke_pid(pid)
1177
1178
1179def is_virtual_machine():
1180    if 'QEMU' in platform.processor():
1181        return True
1182
1183    try:
1184        with open('/sys/devices/virtual/dmi/id/sys_vendor') as f:
1185            if 'QEMU' in f.read():
1186                return True
1187    except IOError:
1188        pass
1189
1190    return False
1191
1192
1193def save_vm_state(checkpoint):
1194    """Saves the current state of the virtual machine.
1195
1196    This function is a NOOP if the test is not running under a virtual machine
1197    with the USB serial port redirected.
1198
1199    Arguments:
1200      checkpoint - Name used to identify this state
1201
1202    Returns:
1203      None
1204    """
1205    # The QEMU monitor has been redirected to the guest serial port located at
1206    # /dev/ttyUSB0. To save the state of the VM, we just send the 'savevm'
1207    # command to the serial port.
1208    if is_virtual_machine() and os.path.exists('/dev/ttyUSB0'):
1209        logging.info('Saving VM state "%s"', checkpoint)
1210        serial = open('/dev/ttyUSB0', 'w')
1211        serial.write('savevm %s\r\n' % checkpoint)
1212        logging.info('Done saving VM state "%s"', checkpoint)
1213
1214
1215def mounts():
1216    ret = []
1217    with open('/proc/mounts') as f:
1218        lines = f.readlines()
1219    for line in lines:
1220        m = re.match(
1221            r'(?P<src>\S+) (?P<dest>\S+) (?P<type>\S+) (?P<opts>\S+).*', line)
1222        if m:
1223            ret.append(m.groupdict())
1224    return ret
1225
1226
1227def is_mountpoint(path):
1228    return path in [m['dest'] for m in mounts()]
1229
1230
1231def require_mountpoint(path):
1232    """
1233    Raises an exception if path is not a mountpoint.
1234    """
1235    if not is_mountpoint(path):
1236        raise error.TestFail('Path not mounted: "%s"' % path)
1237
1238
1239def random_username():
1240    return str(uuid.uuid4()) + '@example.com'
1241
1242
1243def get_signin_credentials(filepath):
1244    """Returns user_id, password tuple from credentials file at filepath.
1245
1246    File must have one line of the format user_id:password
1247
1248    @param filepath: path of credentials file.
1249    @return user_id, password tuple.
1250    """
1251    user_id, password = None, None
1252    if os.path.isfile(filepath):
1253        with open(filepath) as f:
1254            user_id, password = f.read().rstrip().split(':')
1255    return user_id, password
1256
1257
1258def parse_cmd_output(command, run_method=utils.run):
1259    """Runs a command on a host object to retrieve host attributes.
1260
1261    The command should output to stdout in the format of:
1262    <key> = <value> # <optional_comment>
1263
1264
1265    @param command: Command to execute on the host.
1266    @param run_method: Function to use to execute the command. Defaults to
1267                       utils.run so that the command will be executed locally.
1268                       Can be replace with a host.run call so that it will
1269                       execute on a DUT or external machine. Method must accept
1270                       a command argument, stdout_tee and stderr_tee args and
1271                       return a result object with a string attribute stdout
1272                       which will be parsed.
1273
1274    @returns a dictionary mapping host attributes to their values.
1275    """
1276    result = {}
1277    # Suppresses stdout so that the files are not printed to the logs.
1278    cmd_result = run_method(command, stdout_tee=None, stderr_tee=None)
1279    for line in cmd_result.stdout.splitlines():
1280        # Lines are of the format "<key>     = <value>      # <comment>"
1281        key_value = re.match(r'^\s*(?P<key>[^ ]+)\s*=\s*(?P<value>[^ '
1282                             r']+)(?:\s*#.*)?$', line)
1283        if key_value:
1284            result[key_value.group('key')] = key_value.group('value')
1285    return result
1286
1287
1288def set_from_keyval_output(out, delimiter=' '):
1289    """Parse delimiter-separated key-val output into a set of tuples.
1290
1291    Output is expected to be multiline text output from a command.
1292    Stuffs the key-vals into tuples in a set to be later compared.
1293
1294    e.g.  deactivated 0
1295          disableForceClear 0
1296          ==>  set(('deactivated', '0'), ('disableForceClear', '0'))
1297
1298    @param out: multiple lines of space-separated key-val pairs.
1299    @param delimiter: character that separates key from val. Usually a
1300                      space but may be '=' or something else.
1301    @return set of key-val tuples.
1302    """
1303    results = set()
1304    kv_match_re = re.compile('([^ ]+)%s(.*)' % delimiter)
1305    for linecr in out.splitlines():
1306        match = kv_match_re.match(linecr.strip())
1307        if match:
1308            results.add((match.group(1), match.group(2)))
1309    return results
1310
1311
1312def get_cpu_usage():
1313    """Returns machine's CPU usage.
1314
1315    This function uses /proc/stat to identify CPU usage.
1316    Returns:
1317        A dictionary with values for all columns in /proc/stat
1318        Sample dictionary:
1319        {
1320            'user': 254544,
1321            'nice': 9,
1322            'system': 254768,
1323            'idle': 2859878,
1324            'iowait': 1,
1325            'irq': 2,
1326            'softirq': 3,
1327            'steal': 4,
1328            'guest': 5,
1329            'guest_nice': 6
1330        }
1331        If a column is missing or malformed in /proc/stat (typically on older
1332        systems), the value for that column is set to 0.
1333    """
1334    with _open_file('/proc/stat') as proc_stat:
1335        cpu_usage_str = proc_stat.readline().split()
1336    columns = ('user', 'nice', 'system', 'idle', 'iowait', 'irq', 'softirq',
1337               'steal', 'guest', 'guest_nice')
1338    d = {}
1339    for index, col in enumerate(columns, 1):
1340        try:
1341            d[col] = int(cpu_usage_str[index])
1342        except:
1343            d[col] = 0
1344    return d
1345
1346def compute_active_cpu_time(cpu_usage_start, cpu_usage_end):
1347    """Computes the fraction of CPU time spent non-idling.
1348
1349    This function should be invoked using before/after values from calls to
1350    get_cpu_usage().
1351
1352    See https://stackoverflow.com/a/23376195 and
1353    https://unix.stackexchange.com/a/303224 for some more context how
1354    to calculate usage given two /proc/stat snapshots.
1355    """
1356    idle_cols = ('idle', 'iowait')  # All other cols are calculated as active.
1357    time_active_start = sum([x[1] for x in six.iteritems(cpu_usage_start)
1358                             if x[0] not in idle_cols])
1359    time_active_end = sum([x[1] for x in six.iteritems(cpu_usage_end)
1360                           if x[0] not in idle_cols])
1361    total_time_start = sum(cpu_usage_start.values())
1362    total_time_end = sum(cpu_usage_end.values())
1363    # Avoid bogus division which has been observed on Tegra.
1364    if total_time_end <= total_time_start:
1365        logging.warning('compute_active_cpu_time observed bogus data')
1366        # We pretend to be busy, this will force a longer wait for idle CPU.
1367        return 1.0
1368    return ((float(time_active_end) - time_active_start) /
1369            (total_time_end - total_time_start))
1370
1371
1372def is_pgo_mode():
1373    return 'USE_PGO' in os.environ
1374
1375
1376def wait_for_idle_cpu(timeout, utilization):
1377    """Waits for the CPU to become idle (< utilization).
1378
1379    Args:
1380        timeout: The longest time in seconds to wait before throwing an error.
1381        utilization: The CPU usage below which the system should be considered
1382                idle (between 0 and 1.0 independent of cores/hyperthreads).
1383    """
1384    time_passed = 0.0
1385    fraction_active_time = 1.0
1386    sleep_time = 1
1387    logging.info('Starting to wait up to %.1fs for idle CPU...', timeout)
1388    while fraction_active_time >= utilization:
1389        cpu_usage_start = get_cpu_usage()
1390        # Split timeout interval into not too many chunks to limit log spew.
1391        # Start at 1 second, increase exponentially
1392        time.sleep(sleep_time)
1393        time_passed += sleep_time
1394        sleep_time = min(16.0, 2.0 * sleep_time)
1395        cpu_usage_end = get_cpu_usage()
1396        fraction_active_time = compute_active_cpu_time(cpu_usage_start,
1397                                                       cpu_usage_end)
1398        logging.info('After waiting %.1fs CPU utilization is %.3f.',
1399                     time_passed, fraction_active_time)
1400        if time_passed > timeout:
1401            logging.warning('CPU did not become idle.')
1402            log_process_activity()
1403            # crosbug.com/37389
1404            if is_pgo_mode():
1405                logging.info('Still continuing because we are in PGO mode.')
1406                return True
1407
1408            return False
1409    logging.info('Wait for idle CPU took %.1fs (utilization = %.3f).',
1410                 time_passed, fraction_active_time)
1411    return True
1412
1413
1414def log_process_activity():
1415    """Logs the output of top.
1416
1417    Useful to debug performance tests and to find runaway processes.
1418    """
1419    logging.info('Logging current process activity using top and ps.')
1420    cmd = 'top -b -n1 -c'
1421    output = utils.run(cmd)
1422    logging.info(output)
1423    output = utils.run('ps axl')
1424    logging.info(output)
1425
1426
1427def wait_for_cool_machine():
1428    """
1429    A simple heuristic to wait for a machine to cool.
1430    The code looks a bit 'magic', but we don't know ambient temperature
1431    nor machine characteristics and still would like to return the caller
1432    a machine that cooled down as much as reasonably possible.
1433    """
1434    temperature = get_current_temperature_max()
1435    # We got here with a cold machine, return immediately. This should be the
1436    # most common case.
1437    if temperature < 45:
1438        return True
1439    logging.info('Got a hot machine of %dC. Sleeping 1 minute.', temperature)
1440    # A modest wait should cool the machine.
1441    time.sleep(60.0)
1442    temperature = get_current_temperature_max()
1443    # Atoms idle below 60 and everyone else should be even lower.
1444    if temperature < 62:
1445        return True
1446    # This should be rare.
1447    logging.info('Did not cool down (%dC). Sleeping 2 minutes.', temperature)
1448    time.sleep(120.0)
1449    temperature = get_current_temperature_max()
1450    # A temperature over 65'C doesn't give us much headroom to the critical
1451    # temperatures that start at 85'C (and PerfControl as of today will fail at
1452    # critical - 10'C).
1453    if temperature < 65:
1454        return True
1455    logging.warning('Did not cool down (%dC), giving up.', temperature)
1456    log_process_activity()
1457    return False
1458
1459
1460def report_temperature(test, keyname):
1461    """Report current max observed temperature with given keyname.
1462
1463    @param test: autotest_lib.client.bin.test.test instance
1464    @param keyname: key to be used when reporting perf value.
1465    """
1466    temperature = get_current_temperature_max()
1467    logging.info('%s = %f degree Celsius', keyname, temperature)
1468    test.output_perf_value(
1469        description=keyname,
1470        value=temperature,
1471        units='Celsius',
1472        higher_is_better=False)
1473
1474
1475# System paths for machine performance state.
1476_CPUINFO = '/proc/cpuinfo'
1477_DIRTY_WRITEBACK_CENTISECS = '/proc/sys/vm/dirty_writeback_centisecs'
1478_KERNEL_MAX = '/sys/devices/system/cpu/kernel_max'
1479_MEMINFO = '/proc/meminfo'
1480_TEMP_SENSOR_RE = 'Reading temperature...([0-9]*)'
1481
1482def _open_file(path):
1483    """
1484    Opens a file and returns the file object.
1485
1486    This method is intended to be mocked by tests.
1487    @return The open file object.
1488    """
1489    return open(path)
1490
1491def _get_line_from_file(path, line):
1492    """
1493    line can be an integer or
1494    line can be a string that matches the beginning of the line
1495    """
1496    with _open_file(path) as f:
1497        if isinstance(line, int):
1498            l = f.readline()
1499            for _ in range(0, line):
1500                l = f.readline()
1501            return l
1502        else:
1503            for l in f:
1504                if l.startswith(line):
1505                    return l
1506    return None
1507
1508
1509def _get_match_from_file(path, line, prefix, postfix):
1510    """
1511    Matches line in path and returns string between first prefix and postfix.
1512    """
1513    match = _get_line_from_file(path, line)
1514    # Strip everything from front of line including prefix.
1515    if prefix:
1516        match = re.split(prefix, match)[1]
1517    # Strip everything from back of string including first occurence of postfix.
1518    if postfix:
1519        match = re.split(postfix, match)[0]
1520    return match
1521
1522
1523def _get_float_from_file(path, line, prefix, postfix):
1524    match = _get_match_from_file(path, line, prefix, postfix)
1525    return float(match)
1526
1527
1528def _get_int_from_file(path, line, prefix, postfix):
1529    match = _get_match_from_file(path, line, prefix, postfix)
1530    return int(match)
1531
1532
1533def _get_hex_from_file(path, line, prefix, postfix):
1534    match = _get_match_from_file(path, line, prefix, postfix)
1535    return int(match, 16)
1536
1537
1538def is_system_thermally_throttled():
1539    """
1540    Returns whether the system appears to be thermally throttled.
1541    """
1542    for path in glob.glob('/sys/class/thermal/cooling_device*/type'):
1543        with _open_file(path) as f:
1544            cdev_type = f.read().strip()
1545
1546        if not (cdev_type == 'Processor' or
1547                cdev_type.startswith('thermal-devfreq') or
1548                cdev_type.startswith('thermal-cpufreq')):
1549            continue
1550
1551        cur_state_path = os.path.join(os.path.dirname(path), 'cur_state')
1552        if _get_int_from_file(cur_state_path, 0, None, None) > 0:
1553            return True
1554
1555    return False
1556
1557
1558# The paths don't change. Avoid running find all the time.
1559_hwmon_paths = {}
1560
1561def _get_hwmon_datas(file_pattern):
1562    """Returns a list of reading from hwmon."""
1563    # Some systems like daisy_spring only have the virtual hwmon.
1564    # And other systems like rambi only have coretemp.0. See crbug.com/360249.
1565    #    /sys/class/hwmon/hwmon*/
1566    #    /sys/devices/virtual/hwmon/hwmon*/
1567    #    /sys/devices/platform/coretemp.0/
1568    if file_pattern not in _hwmon_paths:
1569        cmd = 'find /sys/class /sys/devices -name "' + file_pattern + '"'
1570        _hwmon_paths[file_pattern] = \
1571            utils.run(cmd, verbose=False).stdout.splitlines()
1572    for _hwmon_path in _hwmon_paths[file_pattern]:
1573        try:
1574            yield _get_float_from_file(_hwmon_path, 0, None, None) * 0.001
1575        except IOError as err:
1576            # Files under /sys may get truncated and result in ENODATA.
1577            # Ignore those.
1578            if err.errno is not errno.ENODATA:
1579                raise
1580
1581
1582def _get_hwmon_temperatures():
1583    """
1584    Returns the currently observed temperatures from hwmon
1585    """
1586    return list(_get_hwmon_datas('temp*_input'))
1587
1588
1589def _get_thermal_zone_temperatures():
1590    """
1591    Returns the maximum currently observered temperature in thermal_zones.
1592    """
1593    temperatures = []
1594    for path in glob.glob('/sys/class/thermal/thermal_zone*/temp'):
1595        try:
1596            temperatures.append(
1597                _get_float_from_file(path, 0, None, None) * 0.001)
1598        except IOError:
1599            # Some devices (e.g. Veyron) may have reserved thermal zones that
1600            # are not active. Trying to read the temperature value would cause a
1601            # EINVAL IO error.
1602            continue
1603    return temperatures
1604
1605
1606def get_ec_temperatures():
1607    """
1608    Uses ectool to return a list of all sensor temperatures in Celsius.
1609
1610    Output from ectool is either '0: 300' or '0: 300 K' (newer ectool
1611    includes the unit).
1612    """
1613    temperatures = []
1614    try:
1615        full_cmd = 'ectool temps all'
1616        lines = utils.run(full_cmd, verbose=False).stdout.splitlines()
1617        pattern = re.compile('.*: (\d+)')
1618        for line in lines:
1619            matched = pattern.match(line)
1620            temperature = int(matched.group(1)) - 273
1621            temperatures.append(temperature)
1622    except Exception as e:
1623        logging.warning('Unable to read temperature sensors using ectool %s.',
1624                        e)
1625    # Sanity check for real world values.
1626    if not all(10.0 <= temperature <= 150.0 for temperature in temperatures):
1627        logging.warning('Unreasonable EC temperatures: %s.', temperatures)
1628    return temperatures
1629
1630
1631def get_current_temperature_max():
1632    """
1633    Returns the highest reported board temperature (all sensors) in Celsius.
1634    """
1635    all_temps = (_get_hwmon_temperatures() +
1636                 _get_thermal_zone_temperatures() +
1637                 get_ec_temperatures())
1638    if all_temps:
1639        temperature = max(all_temps)
1640    else:
1641        temperature = -1
1642    # Sanity check for real world values.
1643    assert ((temperature > 10.0) and
1644            (temperature < 150.0)), ('Unreasonable temperature %.1fC.' %
1645                                     temperature)
1646    return temperature
1647
1648
1649def get_cpu_max_frequency():
1650    """
1651    Returns the largest of the max CPU core frequencies. The unit is Hz.
1652    """
1653    max_frequency = -1
1654    paths = utils._get_cpufreq_paths('cpuinfo_max_freq')
1655    if not paths:
1656        raise ValueError('Could not find max freq; is cpufreq supported?')
1657    for path in paths:
1658        try:
1659            # Convert from kHz to Hz.
1660            frequency = 1000 * _get_float_from_file(path, 0, None, None)
1661        # CPUs may come and go. A missing entry or two aren't critical.
1662        except IOError:
1663            continue
1664        max_frequency = max(frequency, max_frequency)
1665    # Sanity check.
1666    assert max_frequency > 1e8, ('Unreasonably low CPU frequency: %.1f' %
1667            max_frequency)
1668    return max_frequency
1669
1670
1671def get_board_property(key):
1672    """
1673    Get a specific property from /etc/lsb-release.
1674
1675    @param key: board property to return value for
1676
1677    @return the value or '' if not present
1678    """
1679    with open('/etc/lsb-release') as f:
1680        pattern = '%s=(.*)' % key
1681        pat = re.search(pattern, f.read())
1682        if pat:
1683            return pat.group(1)
1684    return ''
1685
1686
1687def get_board():
1688    """
1689    Get the ChromeOS release board name from /etc/lsb-release.
1690    """
1691    return get_board_property('BOARD')
1692
1693
1694def get_board_type():
1695    """
1696    Get the ChromeOS board type from /etc/lsb-release.
1697
1698    @return device type.
1699    """
1700    return get_board_property('DEVICETYPE')
1701
1702
1703def get_chromeos_version():
1704    """
1705    Get the ChromeOS build version from /etc/lsb-release.
1706
1707    @return chromeos release version.
1708    """
1709    return get_board_property('CHROMEOS_RELEASE_VERSION')
1710
1711
1712def get_platform():
1713    """
1714    Get the ChromeOS platform name.
1715
1716    For unibuild this should be equal to model name.  For non-unibuild
1717    it will either be board name or empty string.  In the case of
1718    empty string return board name to match equivalent logic in
1719    server/hosts/cros_host.py
1720
1721    @returns platform name
1722    """
1723    platform = cros_config.call_cros_config_get_output('/ name', utils.run)
1724    if platform == '':
1725        platform = get_board()
1726    return platform
1727
1728
1729def get_sku():
1730    """
1731    Get the SKU number.
1732
1733    @returns SKU number
1734    """
1735    return cros_config.call_cros_config_get_output('/identity sku-id',
1736                                                   utils.run)
1737
1738
1739def get_ec_version():
1740    """Get the ec version as strings.
1741
1742    @returns a string representing this host's ec version.
1743    """
1744    command = 'mosys ec info -s fw_version'
1745    result = utils.run(command, ignore_status=True)
1746    if result.exit_status != 0:
1747        return ''
1748    return result.stdout.strip()
1749
1750
1751def get_firmware_version():
1752    """Get the firmware version as strings.
1753
1754    @returns a string representing this host's firmware version.
1755    """
1756    return utils.run('crossystem fwid').stdout.strip()
1757
1758
1759def get_hardware_revision():
1760    """Get the hardware revision as strings.
1761
1762    @returns a string representing this host's hardware revision.
1763    """
1764    command = 'mosys platform version'
1765    result = utils.run(command, ignore_status=True)
1766    if result.exit_status != 0:
1767        return ''
1768    return result.stdout.strip()
1769
1770
1771def get_kernel_version():
1772    """Get the kernel version as strings.
1773
1774    @returns a string representing this host's kernel version.
1775    """
1776    return utils.run('uname -r').stdout.strip()
1777
1778
1779def get_cpu_name():
1780    """Get the cpu name as strings.
1781
1782    @returns a string representing this host's cpu name.
1783    """
1784
1785    # Try get cpu name from device tree first
1786    if os.path.exists("/proc/device-tree/compatible"):
1787        command = "sed -e 's/\\x0/\\n/g' /proc/device-tree/compatible | tail -1"
1788        return utils.run(command).stdout.strip().replace(',', ' ')
1789
1790
1791    # Get cpu name from uname -p
1792    command = "uname -p"
1793    ret = utils.run(command).stdout.strip()
1794
1795    # 'uname -p' return variant of unknown or amd64 or x86_64 or i686
1796    # Try get cpu name from /proc/cpuinfo instead
1797    if re.match("unknown|amd64|[ix][0-9]?86(_64)?", ret, re.IGNORECASE):
1798        command = "grep model.name /proc/cpuinfo | cut -f 2 -d: | head -1"
1799        ret = utils.run(command).stdout.strip()
1800
1801    # Remove bloat from CPU name, for example
1802    # 'Intel(R) Core(TM) i5-7Y57 CPU @ 1.20GHz'       -> 'Intel Core i5-7Y57'
1803    # 'Intel(R) Xeon(R) CPU E5-2690 v4 @ 2.60GHz'     -> 'Intel Xeon E5-2690 v4'
1804    # 'AMD A10-7850K APU with Radeon(TM) R7 Graphics' -> 'AMD A10-7850K'
1805    # 'AMD GX-212JC SOC with Radeon(TM) R2E Graphics' -> 'AMD GX-212JC'
1806    trim_re = " (@|processor|apu|soc|radeon).*|\(.*?\)| cpu"
1807    return re.sub(trim_re, '', ret, flags=re.IGNORECASE)
1808
1809
1810def get_screen_resolution():
1811    """Get the screen(s) resolution as strings.
1812    In case of more than 1 monitor, return resolution for each monitor separate
1813    with plus sign.
1814
1815    @returns a string representing this host's screen(s) resolution.
1816    """
1817    command = 'for f in /sys/class/drm/*/*/modes; do head -1 $f; done'
1818    ret = utils.run(command, ignore_status=True)
1819    # We might have Chromebox without a screen
1820    if ret.exit_status != 0:
1821        return ''
1822    return ret.stdout.strip().replace('\n', '+')
1823
1824
1825def get_board_with_frequency_and_memory():
1826    """
1827    Returns a board name modified with CPU frequency and memory size to
1828    differentiate between different board variants. For instance
1829    link -> link_1.8GHz_4GB.
1830    """
1831    board_name = get_board()
1832    if is_virtual_machine():
1833        board = '%s_VM' % board_name
1834    else:
1835        memory = get_mem_total_gb()
1836        # Convert frequency to GHz with 1 digit accuracy after the
1837        # decimal point.
1838        frequency = int(round(get_cpu_max_frequency() * 1e-8)) * 0.1
1839        board = '%s_%1.1fGHz_%dGB' % (board_name, frequency, memory)
1840    return board
1841
1842
1843def get_mem_total():
1844    """
1845    Returns the total memory available in the system in MBytes.
1846    """
1847    mem_total = _get_float_from_file(_MEMINFO, 'MemTotal:', 'MemTotal:', ' kB')
1848    # Sanity check, all Chromebooks have at least 1GB of memory.
1849    assert mem_total > 256 * 1024, 'Unreasonable amount of memory.'
1850    return int(mem_total / 1024)
1851
1852
1853def get_mem_total_gb():
1854    """
1855    Returns the total memory available in the system in GBytes.
1856    """
1857    return int(round(get_mem_total() / 1024.0))
1858
1859
1860def get_mem_free():
1861    """
1862    Returns the currently free memory in the system in MBytes.
1863    """
1864    mem_free = _get_float_from_file(_MEMINFO, 'MemFree:', 'MemFree:', ' kB')
1865    return int(mem_free / 1024)
1866
1867def get_mem_free_plus_buffers_and_cached():
1868    """
1869    Returns the free memory in MBytes, counting buffers and cached as free.
1870
1871    This is most often the most interesting number since buffers and cached
1872    memory can be reclaimed on demand. Note however, that there are cases
1873    where this as misleading as well, for example used tmpfs space
1874    count as Cached but can not be reclaimed on demand.
1875    See https://www.kernel.org/doc/Documentation/filesystems/tmpfs.txt.
1876    """
1877    free_mb = get_mem_free()
1878    cached_mb = (_get_float_from_file(
1879        _MEMINFO, 'Cached:', 'Cached:', ' kB') / 1024)
1880    buffers_mb = (_get_float_from_file(
1881        _MEMINFO, 'Buffers:', 'Buffers:', ' kB') / 1024)
1882    return free_mb + buffers_mb + cached_mb
1883
1884
1885def get_dirty_writeback_centisecs():
1886    """
1887    Reads /proc/sys/vm/dirty_writeback_centisecs.
1888    """
1889    time = _get_int_from_file(_DIRTY_WRITEBACK_CENTISECS, 0, None, None)
1890    return time
1891
1892
1893def set_dirty_writeback_centisecs(time=60000):
1894    """
1895    In hundredths of a second, this is how often pdflush wakes up to write data
1896    to disk. The default wakes up the two (or more) active threads every five
1897    seconds. The ChromeOS default is 10 minutes.
1898
1899    We use this to set as low as 1 second to flush error messages in system
1900    logs earlier to disk.
1901    """
1902    # Flush buffers first to make this function synchronous.
1903    utils.system('sync')
1904    if time >= 0:
1905        cmd = 'echo %d > %s' % (time, _DIRTY_WRITEBACK_CENTISECS)
1906        utils.system(cmd)
1907
1908
1909def wflinfo_cmd():
1910    """
1911    Returns a wflinfo command appropriate to the current graphics platform/api.
1912    """
1913    return 'wflinfo -p %s -a %s' % (graphics_platform(), graphics_api())
1914
1915
1916def has_mali():
1917    """ @return: True if system has a Mali GPU enabled."""
1918    return os.path.exists('/dev/mali0')
1919
1920def get_gpu_family():
1921    """Returns the GPU family name."""
1922    global pciid_to_amd_architecture
1923    global pciid_to_intel_architecture
1924
1925    socfamily = get_cpu_soc_family()
1926    if socfamily == 'exynos5' or socfamily == 'rockchip' or has_mali():
1927        cmd = wflinfo_cmd()
1928        wflinfo = utils.system_output(cmd,
1929                                      retain_output=True,
1930                                      ignore_status=False)
1931        m = re.findall(r'OpenGL renderer string: (Mali-\w+)', wflinfo)
1932        if m:
1933            return m[0].lower()
1934        return 'mali-unrecognized'
1935    if socfamily == 'tegra':
1936        return 'tegra'
1937    if socfamily == 'qualcomm':
1938        return 'qualcomm'
1939    if os.path.exists('/sys/kernel/debug/pvr'):
1940        return 'rogue'
1941
1942    pci_vga_device = utils.run("lspci | grep VGA").stdout.rstrip('\n')
1943    bus_device_function = pci_vga_device.partition(' ')[0]
1944    pci_path = '/sys/bus/pci/devices/0000:' + bus_device_function + '/device'
1945
1946    if not os.path.exists(pci_path):
1947        raise error.TestError('PCI device 0000:' + bus_device_function + ' not found')
1948
1949    device_id = utils.read_one_line(pci_path).lower()
1950
1951    if "Advanced Micro Devices" in pci_vga_device:
1952        if not pciid_to_amd_architecture:
1953            with open(_AMD_PCI_IDS_FILE_PATH, 'r') as in_f:
1954                pciid_to_amd_architecture = json.load(in_f)
1955
1956        return pciid_to_amd_architecture[device_id]
1957
1958    if "Intel Corporation" in pci_vga_device:
1959        # Only load Intel PCI ID file once and only if necessary.
1960        if not pciid_to_intel_architecture:
1961            with open(_INTEL_PCI_IDS_FILE_PATH, 'r') as in_f:
1962                pciid_to_intel_architecture = json.load(in_f)
1963
1964        return pciid_to_intel_architecture[device_id]
1965
1966# TODO(ihf): Consider using /etc/lsb-release DEVICETYPE != CHROMEBOOK/CHROMEBASE
1967# for sanity check, but usage seems a bit inconsistent. See
1968# src/third_party/chromiumos-overlay/eclass/appid.eclass
1969_BOARDS_WITHOUT_MONITOR = [
1970    'anglar', 'mccloud', 'monroe', 'ninja', 'rikku', 'guado', 'jecht', 'tidus',
1971    'beltino', 'panther', 'stumpy', 'panther', 'tricky', 'zako', 'veyron_rialto'
1972]
1973
1974
1975def has_no_monitor():
1976    """Returns whether a machine doesn't have a built-in monitor."""
1977    board_name = get_board()
1978    if board_name in _BOARDS_WITHOUT_MONITOR:
1979        return True
1980
1981    return False
1982
1983
1984def get_fixed_dst_drive():
1985    """
1986    Return device name for internal disk.
1987    Example: return /dev/sda for falco booted from usb
1988    """
1989    cmd = ' '.join(['. /usr/sbin/write_gpt.sh;',
1990                    '. /usr/share/misc/chromeos-common.sh;',
1991                    'load_base_vars;',
1992                    'get_fixed_dst_drive'])
1993    return utils.system_output(cmd)
1994
1995
1996def get_root_device():
1997    """
1998    Return root device.
1999    Will return correct disk device even system boot from /dev/dm-0
2000    Example: return /dev/sdb for falco booted from usb
2001    """
2002    return utils.system_output('rootdev -s -d')
2003
2004
2005def get_other_device():
2006    """
2007    Return the non root devices.
2008    Will return a list of other block devices, that are not the root device.
2009    """
2010
2011    cmd = 'lsblk -dpn -o NAME | grep -v -E "(loop|zram|boot|rpmb)"'
2012    devs = utils.system_output(cmd).splitlines()
2013
2014    for dev in devs[:]:
2015        if not re.match(r'/dev/(sd[a-z]|mmcblk[0-9]+|nvme[0-9]+)p?[0-9]*', dev):
2016            devs.remove(dev)
2017        if dev == get_root_device():
2018            devs.remove(dev)
2019    return devs
2020
2021
2022def get_root_partition():
2023    """
2024    Return current root partition
2025    Example: return /dev/sdb3 for falco booted from usb
2026    """
2027    return utils.system_output('rootdev -s')
2028
2029
2030def get_free_root_partition(root_part=None):
2031    """
2032    Return currently unused root partion
2033    Example: return /dev/sdb5 for falco booted from usb
2034
2035    @param root_part: cuurent root partition
2036    """
2037    spare_root_map = {'3': '5', '5': '3'}
2038    if not root_part:
2039        root_part = get_root_partition()
2040    return root_part[:-1] + spare_root_map[root_part[-1]]
2041
2042
2043def get_kernel_partition(root_part=None):
2044    """
2045    Return current kernel partition
2046    Example: return /dev/sda2 for falco booted from usb
2047
2048    @param root_part: current root partition
2049    """
2050    if not root_part:
2051         root_part = get_root_partition()
2052    current_kernel_map = {'3': '2', '5': '4'}
2053    return root_part[:-1] + current_kernel_map[root_part[-1]]
2054
2055
2056def get_free_kernel_partition(root_part=None):
2057    """
2058    return currently unused kernel partition
2059    Example: return /dev/sda4 for falco booted from usb
2060
2061    @param root_part: current root partition
2062    """
2063    kernel_part = get_kernel_partition(root_part)
2064    spare_kernel_map = {'2': '4', '4': '2'}
2065    return kernel_part[:-1] + spare_kernel_map[kernel_part[-1]]
2066
2067
2068def is_booted_from_internal_disk():
2069    """Return True if boot from internal disk. False, otherwise."""
2070    return get_root_device() == get_fixed_dst_drive()
2071
2072
2073def get_ui_use_flags():
2074    """Parses the USE flags as listed in /etc/ui_use_flags.txt.
2075
2076    @return: A list of flag strings found in the ui use flags file.
2077    """
2078    flags = []
2079    for flag in utils.read_file(_UI_USE_FLAGS_FILE_PATH).splitlines():
2080        # Removes everything after the '#'.
2081        flag_before_comment = flag.split('#')[0].strip()
2082        if len(flag_before_comment) != 0:
2083            flags.append(flag_before_comment)
2084
2085    return flags
2086
2087
2088def graphics_platform():
2089    """
2090    Return a string identifying the graphics platform,
2091    e.g. 'glx' or 'x11_egl' or 'gbm'
2092    """
2093    return 'null'
2094
2095
2096def graphics_api():
2097    """Return a string identifying the graphics api, e.g. gl or gles2."""
2098    use_flags = get_ui_use_flags()
2099    if 'opengles' in use_flags:
2100        return 'gles2'
2101    return 'gl'
2102
2103
2104def is_package_installed(package):
2105    """Check if a package is installed already.
2106
2107    @return: True if the package is already installed, otherwise return False.
2108    """
2109    try:
2110        utils.run(_CHECK_PACKAGE_INSTALLED_COMMAND % package)
2111        return True
2112    except error.CmdError:
2113        logging.warn('Package %s is not installed.', package)
2114        return False
2115
2116
2117def is_python_package_installed(package):
2118    """Check if a Python package is installed already.
2119
2120    @return: True if the package is already installed, otherwise return False.
2121    """
2122    try:
2123        __import__(package)
2124        return True
2125    except ImportError:
2126        logging.warn('Python package %s is not installed.', package)
2127        return False
2128
2129
2130def run_sql_cmd(server, user, password, command, database=''):
2131    """Run the given sql command against the specified database.
2132
2133    @param server: Hostname or IP address of the MySQL server.
2134    @param user: User name to log in the MySQL server.
2135    @param password: Password to log in the MySQL server.
2136    @param command: SQL command to run.
2137    @param database: Name of the database to run the command. Default to empty
2138                     for command that does not require specifying database.
2139
2140    @return: The stdout of the command line.
2141    """
2142    cmd = ('mysql -u%s -p%s --host %s %s -e "%s"' %
2143           (user, password, server, database, command))
2144    # Set verbose to False so the command line won't be logged, as it includes
2145    # database credential.
2146    return utils.run(cmd, verbose=False).stdout
2147
2148
2149def strip_non_printable(s):
2150    """Strip non printable characters from string.
2151
2152    @param s: Input string
2153
2154    @return: The input string with only printable characters.
2155    """
2156    return ''.join(x for x in s if x in string.printable)
2157
2158
2159def recursive_func(obj, func, types, sequence_types=(list, tuple, set),
2160                   dict_types=(dict,), fix_num_key=False):
2161    """Apply func to obj recursively.
2162
2163    This function traverses recursively through any sequence-like and
2164    dict-like elements in obj.
2165
2166    @param obj: the object to apply the function func recursively.
2167    @param func: the function to invoke.
2168    @param types: the target types in the object to apply func.
2169    @param sequence_types: the sequence types in python.
2170    @param dict_types: the dict types in python.
2171    @param fix_num_key: to indicate if the key of a dict should be
2172            converted from str type to a number, int or float, type.
2173            It is a culprit of json that it always treats the key of
2174            a dict as string.
2175            Refer to https://docs.python.org/2/library/json.html
2176            for more information.
2177
2178    @return: the result object after applying the func recursively.
2179    """
2180    def ancestors(obj, types):
2181        """Any ancestor of the object class is a subclass of the types?
2182
2183        @param obj: the object to apply the function func.
2184        @param types: the target types of the object.
2185
2186        @return: True if any ancestor class of the obj is found in types;
2187                 False otherwise.
2188        """
2189        return any([issubclass(anc, types) for anc in type(obj).__mro__])
2190
2191    if isinstance(obj, sequence_types) or ancestors(obj, sequence_types):
2192        result_lst = [recursive_func(elm, func, types, fix_num_key=fix_num_key)
2193                      for elm in obj]
2194        # Convert the result list to the object's original sequence type.
2195        return type(obj)(result_lst)
2196    elif isinstance(obj, dict_types) or ancestors(obj, dict_types):
2197        result_lst = [
2198                (recursive_func(key, func, types, fix_num_key=fix_num_key),
2199                 recursive_func(value, func, types, fix_num_key=fix_num_key))
2200                for (key, value) in obj.items()]
2201        # Convert the result list to the object's original dict type.
2202        return type(obj)(result_lst)
2203    # Here are the basic types.
2204    elif isinstance(obj, types) or ancestors(obj, types):
2205        if fix_num_key:
2206            # Check if this is a int or float
2207            try:
2208                result_obj = int(obj)
2209                return result_obj
2210            except ValueError:
2211                try:
2212                    result_obj = float(obj)
2213                    return result_obj
2214                except ValueError:
2215                    pass
2216
2217        result_obj = func(obj)
2218        return result_obj
2219    else:
2220        return obj
2221
2222
2223def base64_recursive_encode(obj):
2224    """Apply base64 encode recursively into the obj structure.
2225
2226    Most of the string-like types could be traced to basestring and bytearray
2227    as follows:
2228        str: basestring
2229        bytes: basestring
2230        dbus.String: basestring
2231        dbus.Signature: basestring
2232        dbus.ByteArray: basestring
2233
2234    Note that all the above types except dbus.String could be traced back to
2235    str. In order to cover dbus.String, basestring is used as the ancestor
2236    class for string-like types.
2237
2238    The other type that needs encoding with base64 in a structure includes
2239        bytearray: bytearray
2240
2241    The sequence types include (list, tuple, set). The dbus.Array is also
2242    covered as
2243        dbus.Array: list
2244
2245    The base dictionary type is dict. The dbus.Dictionary is also covered as
2246        dbus.Dictionary: dict
2247
2248    An example code and output look like
2249        obj = {'a': 10, 'b': 'hello',
2250               'c': [100, 200, bytearray(b'\xf0\xf1\xf2\xf3\xf4')],
2251               'd': {784: bytearray(b'@\x14\x01P'),
2252                     78.0: bytearray(b'\x10\x05\x0b\x10\xb2\x1b\x00')}}
2253        encode_obj = base64_recursive_encode(obj)
2254        decode_obj = base64_recursive_decode(encode_obj)
2255
2256        print 'obj: ', obj
2257        print 'encode_obj: ', encode_obj
2258        print 'decode_obj: ', decode_obj
2259        print 'Equal?', obj == decode_obj
2260
2261        Output:
2262        obj:  {'a': 10,
2263               'c': [100, 200, bytearray(b'\xf0\xf1\xf2\xf3\xf4')],
2264               'b': 'hello',
2265               'd': {784: bytearray(b'@\x14\x01P'),
2266                     78.0: bytearray(b'\x10\x05\x0b\x10\xb2\x1b\x00')}}
2267
2268        encode_obj:  {'YQ==': 10,
2269                      'Yw==': [100, 200, '8PHy8/Q='],
2270                      'Yg==': 'aGVsbG8='
2271                      'ZA==': {784: 'QBQBUA==', 78.0: 'EAULELIbAA=='}}
2272        decode_obj:  {'a': 10,
2273                      'c': [100, 200, '\xf0\xf1\xf2\xf3\xf4'],
2274                      'b': 'hello',
2275                      'd': {784: '@\x14\x01P',
2276                            78.0: '\x10\x05\x0b\x10\xb2\x1b\x00'}}
2277        Equal? True
2278
2279    @param obj: the object to apply base64 encoding recursively.
2280
2281    @return: the base64 encoded object.
2282    """
2283    encode_types = (six.string_types, bytearray)
2284    return recursive_func(obj, base64.standard_b64encode, encode_types)
2285
2286
2287def base64_recursive_decode(obj):
2288    """Apply base64 decode recursively into the obj structure.
2289
2290    @param obj: the object to apply base64 decoding recursively.
2291
2292    @return: the base64 decoded object.
2293    """
2294    decode_types = (six.string_types,)
2295    return recursive_func(obj, base64.standard_b64decode, decode_types,
2296                          fix_num_key=True)
2297