[Autotest] Add test for DeviceAutoUpdateDisabled policy

Add an AU test which checks the behavior of the
DeviceAutoUpdateDisabled device policy.

Server test wrapper policy_AUServer:
- clears the TPM (both before and after the test to ensure
  no lab machines are left in an enrolled state)
- stages a payload in a reachable location
- kicks off the client test

Client test policy_DeviceAutoUpdateDisabled:
- fake-enrolls the device and sets the policy as desired
- sets up a NanoOmaha server to respond to update requests
- verifies that update request is or is not sent, as per policy
- verifies that update does or does not start, as per policy

There are three control files, each with a different policy value.
When the policy is set to True, update is disabled and no update
requests are sent.  When the policy is set to False or not set at
all (here None is used to represent that), update is not disabled
and device can update.

enterprise_au_context adds some helpful functions that are useful
for policy AU tests but are likely too specific for
update_engine_util.

The enterprise_policy_base changes are to handle the difference
between what the FakeDMS expects and the policy name from the
policy template (which is also what is listed in chrome://policy).
This dict will not be difficult to maintain in the long term.

TEST=ran on several lab machines with no issue
BUG=None

Change-Id: Iee6cf840e09e0705b166636dbf896c46d3382662
Reviewed-on: https://chromium-review.googlesource.com/1144536
Commit-Ready: Katherine Threlkeld <kathrelkeld@chromium.org>
Tested-by: Katherine Threlkeld <kathrelkeld@chromium.org>
Reviewed-by: David Haddock <dhaddock@chromium.org>
diff --git a/client/cros/enterprise/enterprise_au_context.py b/client/cros/enterprise/enterprise_au_context.py
new file mode 100644
index 0000000..35c5ac2
--- /dev/null
+++ b/client/cros/enterprise/enterprise_au_context.py
@@ -0,0 +1,87 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import logging
+
+from autotest_lib.client.bin import utils
+from autotest_lib.client.common_lib import error
+from autotest_lib.client.cros.update_engine import nano_omaha_devserver
+from autotest_lib.client.cros.update_engine import update_engine_util
+
+_MIN_BUILD = '0.0.0'
+_MAX_BUILD = '999999.0.0'
+
+class NanoOmahaEnterpriseAUContext(object):
+    """
+    Contains methods required for Enterprise AU tests using Nano Omaha.
+
+    """
+
+    def __init__(self, image_url, image_size, sha256, to_build=_MAX_BUILD,
+                 from_build=_MIN_BUILD, is_rollback=False, is_critical=False):
+        """
+        Start a Nano Omaha instance and intialize variables.
+
+        @param image_url: Url of update image.
+        @param image_size: Size of the update.
+        @param sha256: Sha256 hash of the update.
+        @param to_build: String of the build number Nano Omaha should serve.
+        @param from_build: String of the build number this device should say
+                           it is on by setting lsb_release.
+        @param is_rollback: whether the build should serve with the rollback
+                            flag.
+        @param is_critical: whether the build should serve marked as critical.
+
+        """
+        self._omaha = nano_omaha_devserver.NanoOmahaDevserver()
+        self._omaha.set_image_params(image_url, image_size, sha256,
+                                     build=to_build, is_rollback=is_rollback)
+        self._omaha.start()
+
+        self._au_util = update_engine_util.UpdateEngineUtil()
+
+        update_url = self._omaha.get_update_url()
+        self._au_util._create_custom_lsb_release(from_build, update_url)
+
+
+    def update_and_poll_for_update_start(self):
+        """
+        Check for an update and wait until it starts.
+
+        @raises: error.TestFail when update does not start after timeout.
+
+        """
+        self._au_util._check_for_update(port=self._omaha.get_port())
+
+        def update_started():
+            """Polling function: True or False if update has started."""
+            status = self._au_util._get_update_engine_status()
+            logging.info('Status: %s', status)
+            return (status[self._au_util._CURRENT_OP]
+                    == self._au_util._UPDATE_ENGINE_DOWNLOADING)
+
+        utils.poll_for_condition(
+                update_started,
+                exception=error.TestFail('Update did not start!'))
+
+
+    def get_update_requests(self):
+        """
+        Get the contents of all the update requests from the most recent log.
+
+        @returns: a sequential list of <request> xml blocks or None if none.
+
+        """
+        return self._au_util._get_update_requests()
+
+
+    def get_time_of_last_update_request(self):
+        """
+        Get the time of the last update request from most recent logfile.
+
+        @returns: seconds since epoch of when last update request happened
+                  (second accuracy), or None if no such timestamp exists.
+
+        """
+        return self._au_util._get_time_of_last_update_request()
diff --git a/client/cros/enterprise/enterprise_policy_base.py b/client/cros/enterprise/enterprise_policy_base.py
index aa32247..b9acf96 100755
--- a/client/cros/enterprise/enterprise_policy_base.py
+++ b/client/cros/enterprise/enterprise_policy_base.py
@@ -55,6 +55,12 @@
 PASSWORD = 'fakepassword'
 GAIA_ID = 'fake-gaia-id'
 
+# Convert from chrome://policy name to what fake dms expects.
+DEVICE_POLICY_DICT = {
+    'DeviceAutoUpdateDisabled': 'update_disabled',
+    'DeviceTargetVersionPrefix': 'target_version_prefix',
+    'DeviceRollbackToTargetVersion': 'rollback_to_target_version'
+}
 
 class EnterprisePolicyTest(test.test):
     """Base class for Enterprise Policy Tests."""
@@ -81,6 +87,9 @@
         """
         Initialize test parameters and fake DM Server.
 
+        This function exists so that ARC++ tests (which inherit from the
+        ArcTest class) can also initialize a policy setup.
+
         @param case: String name of the test case to run.
         @param env: String environment of DMS and Gaia servers.
         @param username: String user name login credential.
@@ -234,10 +243,17 @@
         s_user_p = copy.deepcopy(suggested_user_policies)
         device_p = copy.deepcopy(device_policies)
 
+        # Replace all device policies with their FakeDMS-friendly names.
+        fixed_device_p = {}
+        for policy in device_p:
+            if policy not in DEVICE_POLICY_DICT:
+                raise error.TestError('Cannot convert %s!' % policy)
+            fixed_device_p[DEVICE_POLICY_DICT[policy]] = device_p[policy]
+
         # Remove "Not set" policies and json-ify dicts because the
         # FakeDMServer expects "policy": "{value}" not "policy": {value}
         # and "policy": "[{value}]" not "policy": [{value}].
-        for policies_dict in [user_p, s_user_p, device_p]:
+        for policies_dict in [user_p, s_user_p, fixed_device_p]:
             policies_to_pop = []
             for policy in policies_dict:
                 value = policies_dict[policy]
@@ -267,9 +283,8 @@
                 user_modes_dict['recommended'] = s_user_p
             management_dict['google/chromeos/user'] = user_modes_dict
 
-        if device_p:
-            management_dict['google/chromeos/device'] = device_p
-
+        if fixed_device_p:
+            management_dict['google/chromeos/device'] = fixed_device_p
 
         logging.info('Created policy blob: %s', management_dict)
         return encode_json_string(management_dict)
@@ -532,34 +547,37 @@
         logging.info('  gaia_login: %s', not self.dms_is_fake)
 
         if enroll:
-            self.cr = chrome.Chrome(auto_login=False,
-                                    extra_browser_args=extra_flags,
-                                    expect_policy_fetch=True)
+            self.cr = chrome.Chrome(
+                    auto_login=False,
+                    extra_browser_args=extra_flags,
+                    expect_policy_fetch=True)
             if self.dms_is_fake:
                 enrollment.EnterpriseFakeEnrollment(
-                    self.cr.browser, self.username, self.password, self.gaia_id,
-                    auto_login=auto_login)
+                        self.cr.browser, self.username, self.password,
+                        self.gaia_id, auto_login=auto_login)
             else:
                 enrollment.EnterpriseEnrollment(
-                    self.cr.browser, self.username, self.password,
-                    auto_login=auto_login)
+                        self.cr.browser, self.username, self.password,
+                        auto_login=auto_login)
 
         elif auto_login:
-            self.cr = chrome.Chrome(extra_browser_args=extra_flags,
-                                    username=self.username,
-                                    password=self.password,
-                                    gaia_login=not self.dms_is_fake,
-                                    disable_gaia_services=self.dms_is_fake,
-                                    autotest_ext=True,
-                                    init_network_controller=init_network_controller,
-                                    expect_policy_fetch=True,
-                                    extension_paths=extension_paths)
+            self.cr = chrome.Chrome(
+                    extra_browser_args=extra_flags,
+                    username=self.username,
+                    password=self.password,
+                    gaia_login=not self.dms_is_fake,
+                    disable_gaia_services=self.dms_is_fake,
+                    autotest_ext=True,
+                    init_network_controller=init_network_controller,
+                    expect_policy_fetch=True,
+                    extension_paths=extension_paths)
         else:
-            self.cr = chrome.Chrome(auto_login=False,
-                                    extra_browser_args=extra_flags,
-                                    disable_gaia_services=self.dms_is_fake,
-                                    autotest_ext=True,
-                                    expect_policy_fetch=True)
+            self.cr = chrome.Chrome(
+                    auto_login=False,
+                    extra_browser_args=extra_flags,
+                    disable_gaia_services=self.dms_is_fake,
+                    autotest_ext=True,
+                    expect_policy_fetch=True)
 
         if auto_login:
             if not cryptohome.is_vault_mounted(user=self.username,
diff --git a/client/cros/update_engine/update_engine_util.py b/client/cros/update_engine/update_engine_util.py
index 5df5e1b..c5da050 100644
--- a/client/cros/update_engine/update_engine_util.py
+++ b/client/cros/update_engine/update_engine_util.py
@@ -2,8 +2,10 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import datetime
 import logging
 import os
+import re
 import shutil
 import time
 
@@ -318,6 +320,52 @@
         self._run('rm %s' % self._CUSTOM_LSB_RELEASE, ignore_status=True)
 
 
+    def _get_update_requests(self):
+        """
+        Get the contents of all the update requests from the most recent log.
+
+        @returns: a sequential list of <request> xml blocks or None if none.
+
+        """
+        update_log = ''
+        with open(self._UPDATE_ENGINE_LOG) as fh:
+            update_log = fh.read()
+
+        # Matches <request> ... </request>.  The match can be on multiple
+        # lines and the search is not greedy so it only matches one block.
+        return re.findall(r'<request>.?</request>', update_log, re.DOTALL)
+
+
+    def _get_time_of_last_update_request(self):
+        """
+        Get the time of the last update request from most recent logfile.
+
+        @returns: seconds since epoch of when last update request happened
+                  (second accuracy), or None if no such timestamp exists.
+
+        """
+        update_log = ''
+        with open(self._UPDATE_ENGINE_LOG) as fh:
+            update_log = fh.read()
+
+        # Matches any single line with "MMDD/HHMMSS ... Request ... xml", e.g.
+        # "[0723/133526:INFO:omaha_request_action.cc(794)] Request: <?xml".
+        result = re.findall(r'([0-9]{4}/[0-9]{6}).*Request.*xml', update_log)
+        if not result:
+            return None
+
+        LOG_TIMESTAMP_FORMAT = '%m%d/%H%M%S'
+        match = result[-1]
+
+        # The log does not include the year, so set it as this year.
+        # This assumption could cause incorrect behavior, but is unlikely to.
+        current_year = datetime.datetime.now().year
+        log_datetime = datetime.datetime.strptime(match, LOG_TIMESTAMP_FORMAT)
+        log_datetime = log_datetime.replace(year=current_year)
+
+        return time.mktime(log_datetime.timetuple())
+
+
     def _take_screenshot(self, filename):
         """
         Take a screenshot and save in resultsdir.
diff --git a/client/site_tests/policy_DeviceAutoUpdateDisabled/control b/client/site_tests/policy_DeviceAutoUpdateDisabled/control
new file mode 100644
index 0000000..7122655
--- /dev/null
+++ b/client/site_tests/policy_DeviceAutoUpdateDisabled/control
@@ -0,0 +1,19 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+AUTHOR = 'kathrelkeld'
+NAME = 'policy_DeviceAutoUpdateDisabled'
+TIME = 'SHORT'
+TEST_CATEGORY = 'General'
+TEST_CLASS = 'enterprise'
+TEST_TYPE = 'client'
+
+DOC = '''
+Called through the policy_AUServer test only.  Verifies whether the device
+can or cannot Autoupdate with the DeviceAutoUpdateDisabled policy set.
+'''
+
+args_dict = utils.args_to_dict(args)
+
+job.run_test('policy_DeviceAutoUpdateDisabled', **args_dict)
diff --git a/client/site_tests/policy_DeviceAutoUpdateDisabled/policy_DeviceAutoUpdateDisabled.py b/client/site_tests/policy_DeviceAutoUpdateDisabled/policy_DeviceAutoUpdateDisabled.py
new file mode 100644
index 0000000..dfea1d7
--- /dev/null
+++ b/client/site_tests/policy_DeviceAutoUpdateDisabled/policy_DeviceAutoUpdateDisabled.py
@@ -0,0 +1,70 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import logging
+import math
+import time
+
+from autotest_lib.client.common_lib import error
+from autotest_lib.client.cros.enterprise import enterprise_au_context
+from autotest_lib.client.cros.enterprise import enterprise_policy_base
+
+
+class policy_DeviceAutoUpdateDisabled(
+        enterprise_policy_base.EnterprisePolicyTest):
+    """Test for the DeviceAutoUpdateDisabled policy."""
+    version = 1
+    _POLICY = 'DeviceAutoUpdateDisabled'
+
+
+    def _test_update_disabled(self, should_update):
+        """
+        Main test function.
+
+        Try to update and poll for start (or lack of start) to the update.
+        Check whether an update request was sent.
+
+        @param should_update: True or False whether the device should update.
+
+        """
+        # Log time is only in second accuracy.  Assume no update request has
+        # occured since the current whole second started.
+        start_time = math.floor(time.time())
+        logging.info('Update test start time: %s', start_time)
+
+        try:
+            self._au_context.update_and_poll_for_update_start()
+        except error.TestFail as e:
+            if should_update:
+                raise e
+        else:
+            if not should_update:
+                raise error.TestFail('Update started when it should not have!')
+
+        update_time = self._au_context.get_time_of_last_update_request()
+        logging.info('Last update time: %s', update_time)
+
+        if should_update and (not update_time or update_time < start_time):
+            raise error.TestFail('No update request was sent!')
+        if not should_update and update_time and update_time >= start_time:
+            raise error.TestFail('Update request was sent!')
+
+
+    def run_once(self, case, image_url, image_size, sha256, enroll=True):
+        """
+        Entry point of this test.
+
+        @param case: True, False, or None for the value of the update policy.
+        @param image_url: Url of update image (this build).
+        @param image_size: Size of the update.
+        @param sha256: Sha256 hash of the update.
+
+        """
+        self.setup_case(device_policies={self._POLICY: case}, enroll=enroll)
+
+        self._au_context = enterprise_au_context.NanoOmahaEnterpriseAUContext(
+                image_url=image_url, image_size=image_size, sha256=sha256)
+
+        # When policy is False or not set, user should update.
+        self._test_update_disabled(should_update=case is not True)
diff --git a/server/site_tests/policy_AUServer/control.AUDisabled.false b/server/site_tests/policy_AUServer/control.AUDisabled.false
new file mode 100644
index 0000000..a0af84b
--- /dev/null
+++ b/server/site_tests/policy_AUServer/control.AUDisabled.false
@@ -0,0 +1,29 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+from autotest_lib.server import utils
+
+AUTHOR = 'kathrelkeld'
+NAME = 'policy_AUServer.AUDisabled.false'
+TIME = 'SHORT'
+TEST_CATEGORY = 'General'
+TEST_CLASS = 'enterprise'
+TEST_TYPE = 'server'
+ATTRIBUTES = 'suite:ent-nightly, suite:policy'
+
+DOC = """
+Sets up and runs the client test for the DeviceAutoUpdateDisabled
+policy.
+
+"""
+args_dict = utils.args_to_dict(args)
+client_test = 'policy_DeviceAutoUpdateDisabled'
+case = False
+
+def run(machine):
+    host = hosts.create_host(machine)
+    job.run_test('policy_AUServer', host=host, client_test=client_test,
+                 case=case, **args_dict)
+
+parallel_simple(run, machines)
diff --git a/server/site_tests/policy_AUServer/control.AUDisabled.notset b/server/site_tests/policy_AUServer/control.AUDisabled.notset
new file mode 100644
index 0000000..020468c
--- /dev/null
+++ b/server/site_tests/policy_AUServer/control.AUDisabled.notset
@@ -0,0 +1,29 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+from autotest_lib.server import utils
+
+AUTHOR = 'kathrelkeld'
+NAME = 'policy_AUServer.AUDisabled.notset'
+TIME = 'SHORT'
+TEST_CATEGORY = 'General'
+TEST_CLASS = 'enterprise'
+TEST_TYPE = 'server'
+ATTRIBUTES = 'suite:ent-nightly, suite:policy'
+
+DOC = """
+Sets up and runs the client test for the DeviceAutoUpdateDisabled
+policy.
+
+"""
+args_dict = utils.args_to_dict(args)
+client_test = 'policy_DeviceAutoUpdateDisabled'
+case = None
+
+def run(machine):
+    host = hosts.create_host(machine)
+    job.run_test('policy_AUServer', host=host, client_test=client_test,
+                 case=case, **args_dict)
+
+parallel_simple(run, machines)
diff --git a/server/site_tests/policy_AUServer/control.AUDisabled.true b/server/site_tests/policy_AUServer/control.AUDisabled.true
new file mode 100644
index 0000000..72178f4
--- /dev/null
+++ b/server/site_tests/policy_AUServer/control.AUDisabled.true
@@ -0,0 +1,29 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+from autotest_lib.server import utils
+
+AUTHOR = 'kathrelkeld'
+NAME = 'policy_AUServer.AUDisabled.true'
+TIME = 'SHORT'
+TEST_CATEGORY = 'General'
+TEST_CLASS = 'enterprise'
+TEST_TYPE = 'server'
+ATTRIBUTES = 'suite:ent-nightly, suite:policy'
+
+DOC = """
+Sets up and runs the client test for the DeviceAutoUpdateDisabled
+policy.
+
+"""
+args_dict = utils.args_to_dict(args)
+client_test = 'policy_DeviceAutoUpdateDisabled'
+case = True
+
+def run(machine):
+    host = hosts.create_host(machine)
+    job.run_test('policy_AUServer', host=host, client_test=client_test,
+                 case=case, **args_dict)
+
+parallel_simple(run, machines)
diff --git a/server/site_tests/policy_AUServer/policy_AUServer.py b/server/site_tests/policy_AUServer/policy_AUServer.py
new file mode 100644
index 0000000..4a85074
--- /dev/null
+++ b/server/site_tests/policy_AUServer/policy_AUServer.py
@@ -0,0 +1,72 @@
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import logging
+
+from autotest_lib.client.common_lib.cros import tpm_utils
+from autotest_lib.server.cros.update_engine import update_engine_test
+
+
+class policy_AUServer(update_engine_test.UpdateEngineTest):
+    """
+    This server test is used just to get the URL of the payload to use. It
+    will then call into a client side test to test different things in
+    the Omaha response.
+    """
+    version = 1
+
+    def clear_tpm_if_owned(self):
+        """Clear the TPM only if device is already owned."""
+        tpm_status = tpm_utils.TPMStatus(self._host)
+        logging.info('TPM status: %s', tpm_status)
+        if tpm_status['Owned']:
+            logging.info('Clearing TPM because this device is owned.')
+            tpm_utils.ClearTPMOwnerRequest(self._host)
+
+
+    def cleanup(self):
+        """Cleanup for this test."""
+        super(policy_AUServer, self).cleanup()
+        self.clear_tpm_if_owned()
+        self._host.reboot()
+
+
+    def run_once(self, client_test, case, full_payload=True,
+                 job_repo_url=None, running_at_desk=False):
+        """
+        Starting point of this test.
+
+        Note: base class sets host as self._host.
+
+        @param client_test: the name of the Client test to run.
+        @param case: the case to run for the given Client test.
+        @param full_payload: whether the update should be full or incremental.
+        @param job_repo_url: url provided at runtime (or passed in locally
+                             when running at workstation).
+        @param running_at_desk: indicates test is run from a workstation.
+
+        """
+        self._job_repo_url = job_repo_url
+
+        # Clear TPM to ensure that client test can enroll device.
+        self.clear_tpm_if_owned()
+
+        # Figure out the payload to use for the current build.
+        payload = self._get_payload_url(full_payload=full_payload)
+        image_url = self._stage_payload_by_uri(payload)
+        file_info = self._get_staged_file_info(image_url)
+
+        if running_at_desk:
+            image_url = self._copy_payload_to_public_bucket(payload)
+            logging.info('We are running from a workstation. Putting URL on a '
+                         'public location: %s', image_url)
+
+        logging.info('url: %s', image_url)
+        logging.info('file_info: %s', file_info)
+
+        self._run_client_test_and_check_result(client_test,
+                                               case=case,
+                                               image_url=image_url,
+                                               image_size=file_info['size'],
+                                               sha256=file_info['sha256'])