1# Copyright 2020 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5from autotest_lib.client.common_lib import error
6from chromite.lib import retry_util
7
8
9_VPD_BASE_CMD = 'vpd -i %s %s %s'
10_RW = 'RW_VPD'
11_RO = 'RO_VPD'
12
13
14def _check_partition(partition):
15    """
16    Used to validate input string in other functions.
17
18    @param partition: If this is not 'RO_VPD' or 'RW_VPD', raise a ValueError.
19
20    """
21    if partition not in [_RW, _RO]:
22        raise ValueError("partition should be 'RW_VPD' or 'RO_VPD'")
23
24
25def dump_vpd_log(host, force=True, retries=3):
26    """
27    Applies changes to the VPD settings by flushing them to the VPD cache and
28    output files.
29
30    @param host: Host to run the command on.
31    @param force: True to pass in the --force parameter to forcefully dump
32                  the log. False to omit it.
33    @param retries: Number of times to try rerunning the command in case of
34                    error.
35
36    """
37    vpd_dump_cmd = 'dump_vpd_log%s' % (' --force' if force else '')
38    retry_util.RetryException(error.AutoservRunError, retries, host.run,
39                              vpd_dump_cmd)
40
41
42def vpd_get(host, key, partition='RW_VPD', retries=3):
43    """
44    Gets the VPD value associated with the input key.
45
46    @param host: Host to run the command on.
47    @param key: Key of the desired VPD value.
48    @param partition: Which partition to access. 'RO_VPD' or 'RW_VPD'.
49    @param retries: Number of times to try rerunning the command in case of
50                    error.
51
52    """
53    _check_partition(partition)
54    get_cmd = _VPD_BASE_CMD % (partition, '-g', key)
55    try:
56        return retry_util.RetryException(error.AutoservRunError, retries,
57                                         host.run, get_cmd).stdout
58    except error.AutoservRunError as e:
59        if 'was not found' in str(e.result_obj.stderr):
60            return None
61        else:
62            raise e
63
64
65def vpd_set(host, vpd_dict, partition='RW_VPD', dump=False, force_dump=False,
66            retries=3):
67    """
68    Sets the given key/value pairs in the specified VPD partition.
69
70    @param host: Host to run the command on.
71    @param vpd_dict: Dictionary containing the VPD key/value pairs to set.
72                     Dictionary keys should be the VPD key strings, and values
73                     should be the desired values to write.
74    @param partition: Which partition to access. 'RO_VPD' or 'RW_VPD'.
75    @param dump: If True, also run dump_vpd_log command after setting the
76                 vpd values.
77    @param force_dump: Whether or not to forcefully dump the vpd log.
78    @param retries: Number of times to try rerunning the command in case of
79                    error.
80
81    """
82    _check_partition(partition)
83    for vpd_key in vpd_dict:
84        set_cmd = _VPD_BASE_CMD % (partition, '-s',
85                  (vpd_key + '=' + str(vpd_dict[vpd_key])))
86        retry_util.RetryException(error.AutoservRunError, retries,
87                                  host.run, set_cmd).stdout
88
89    if dump:
90        dump_vpd_log(host, force=force_dump, retries=retries)
91
92
93def vpd_delete(host, key, partition='RW_VPD', dump=False, force_dump=False,
94               retries=3):
95    """
96    Deletes the specified key from the specified VPD partition.
97
98    @param host: Host to run the command on.
99    @param key: The VPD value to delete.
100    @param partition: Which partition to access. 'RO_VPD' or 'RW_VPD'.
101    @param dump: If True, also run dump_vpd_log command after deleting the
102                 vpd value.
103    @param force_dump: Whether or not to forcefully dump the vpd log.
104    @param retries: Number of times to try rerunning the command in case of
105                    error.
106
107    """
108    _check_partition(partition)
109    if not vpd_get(host, key, partition=partition, retries=retries):
110        return
111
112    del_cmd = _VPD_BASE_CMD % (partition, '-d', key)
113    retry_util.RetryException(error.AutoservRunError, retries, host.run,
114                              del_cmd).stdout
115
116    if dump:
117        dump_vpd_log(host, force=force_dump, retries=retries)
118