1# Copyright 2019 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"""
5
6Helper file for handling policy fetching/parsing/formatting. This file has
7working unitests located in ${THIS_FILE_NAME}_unittest.py. If you modify this
8file, add/fix/adjust the unittests for proper code coverage and
9re-run.
10
11Read instructions in the unittest file on how to run.
12
13"""
14import json
15import time
16
17from autotest_lib.client.common_lib import error
18# Default settings for managed user policies
19
20
21def get_all_policies(autotest_ext):
22    """Returns a dict of all the policies on the device."""
23
24    policy_data = _get_pol_from_api(autotest_ext)
25    if not policy_data:
26        raise error.TestError('API did not return policy data!')
27    _reformat_policies(policy_data)
28
29    return policy_data
30
31
32def _get_pol_from_api(autotest_ext):
33    """Call the getAllPolicies API, return raw result, clear stored data."""
34    new_promise = '''new Promise(function(resolve, reject) {
35        chrome.autotestPrivate.getAllEnterprisePolicies(function(policies) {
36          if (chrome.runtime.lastError) {
37            reject(new Error(chrome.runtime.lastError.message));
38          } else {
39            resolve(policies);
40          }
41        })
42    })'''
43    policy_data = autotest_ext.EvaluateJavaScript(new_promise, promise=True)
44    return policy_data
45
46
47def _wait_for_new_pols_after_refresh(autotest_ext):
48    """
49    Wait up to 1 second for the polices to update when refreshing.
50
51    Note: This is non-breaking, thus even if the policies do not update, once
52    the timeout expires, the function will exit.
53
54    """
55    prior = _get_pol_from_api(autotest_ext)
56    _call_refresh_policies_from_api(autotest_ext)
57    t1 = time.time()
58    while time.time() - t1 < 1:
59        curr = _get_pol_from_api(autotest_ext)
60        if curr != prior:
61            break
62
63
64def _call_refresh_policies_from_api(autotest_ext):
65    """Call the refreshEnterprisePolicies from autotestPrivate."""
66    new_promise = '''new Promise(function(resolve, reject) {
67        chrome.autotestPrivate.refreshEnterprisePolicies(function() {
68          if (chrome.runtime.lastError) {
69            reject(new Error(chrome.runtime.lastError.message));
70          } else {
71            resolve();
72          }
73        })
74    })'''
75    autotest_ext.EvaluateJavaScript(new_promise, promise=True)
76
77
78def refresh_policies(autotest_ext, wait_for_new=False):
79    """Force a policy fetch."""
80
81    if wait_for_new:
82        _wait_for_new_pols_after_refresh(autotest_ext)
83    else:
84        _call_refresh_policies_from_api(autotest_ext)
85
86
87def _reformat_policies(policy_dict):
88    """
89    Reformat visually formatted dicts to type dict (and not unicode).
90
91    Given a dict, check if 'value' is a key in the value field. If so, check
92    if the data of ['value'] is unicode. If it is, attempt to load it as json.
93    If not, but the value field is a dict, recursively call this function to
94    check again.
95
96    @param: policy_dict, a policy dictionary.
97
98    """
99    for k, v in policy_dict.items():
100        if not v:
101            # No data
102            continue
103        if 'value' in v:
104            if type(v['value']) == unicode:
105                _remove_visual_formatting(v)
106        elif isinstance(v, dict):
107            _reformat_policies(v)
108
109
110def _remove_visual_formatting(policy_value):
111    """
112    Attempt to remove any visual formatting.
113
114    Note: policy_value can be either unicode dict or string. If json.loads()
115    raises a ValueError, it is a string, and we do not need to format.
116
117    """
118    try:
119        policy_value['value'] = json.loads(policy_value['value'])
120    except ValueError:
121        return
122