1# Copyright 2018 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
5import logging
6import re
7
8from autotest_lib.client.bin import utils
9from autotest_lib.client.common_lib import error
10from autotest_lib.client.cros.update_engine import nano_omaha_devserver
11from autotest_lib.client.cros.update_engine import update_engine_util
12
13_MIN_BUILD = '1.1.1'
14_MAX_BUILD = '999999.9.9'
15
16# eventtype value sent by client in an AU request.  Indicates that
17# device rebooted after an update since last check.
18# TODO: remove this, crbug.com/879687
19_EVENT_TYPE_REBOOTED_AFTER_UPDATE = '54'
20
21class NanoOmahaEnterpriseAUContext(object):
22    """
23    Contains methods required for Enterprise AU tests using Nano Omaha.
24
25    """
26    def __init__(self, image_url, image_size, sha256, to_build=_MAX_BUILD,
27                 from_build=_MIN_BUILD, is_rollback=False, is_critical=False):
28        """
29        Start a Nano Omaha instance and intialize variables.
30
31        @param image_url: Url of update image.
32        @param image_size: Size of the update.
33        @param sha256: Sha256 hash of the update.
34        @param to_build: String of the build number Nano Omaha should serve.
35        @param from_build: String of the build number this device should say
36                           it is on by setting lsb_release.
37        @param is_rollback: whether the build should serve with the rollback
38                            flag.
39        @param is_critical: whether the build should serve marked as critical.
40
41        """
42        self._omaha = nano_omaha_devserver.NanoOmahaDevserver()
43        self._omaha.set_image_params(image_url, image_size, sha256,
44                                     build=to_build, is_rollback=is_rollback)
45        self._omaha.start()
46
47        self._au_util = update_engine_util.UpdateEngineUtil()
48
49        update_url = self._omaha.get_update_url()
50        self._au_util._create_custom_lsb_release(from_build, update_url)
51
52        self._is_rollback = is_rollback
53        self._is_critical = is_critical
54
55
56    def update_and_poll_for_update_start(self, is_interactive=False):
57        """
58        Check for an update and wait until it starts.
59
60        @param is_interactive: whether the request is interactive.
61
62        @raises: error.TestFail when update does not start after timeout.
63
64        """
65        self._au_util._check_for_update(port=self._omaha.get_port(),
66                                        interactive=is_interactive)
67
68        def update_started():
69            """Polling function: True or False if update has started."""
70            status = self._au_util._get_update_engine_status()
71            logging.info('Status: %s', status)
72            return (status[self._au_util._CURRENT_OP]
73                    == self._au_util._UPDATE_ENGINE_DOWNLOADING)
74
75        utils.poll_for_condition(
76                update_started,
77                exception=error.TestFail('Update did not start!'))
78
79
80    def get_update_requests(self):
81        """
82        Get the contents of all the update requests from the most recent log.
83
84        @returns: a sequential list of <request> xml blocks or None if none.
85
86        """
87        return self._au_util._get_update_requests()
88
89
90    def get_time_of_last_update_request(self):
91        """
92        Get the time of the last update request from most recent logfile.
93
94        @returns: seconds since epoch of when last update request happened
95                  (second accuracy), or None if no such timestamp exists.
96
97        """
98        return self._au_util._get_time_of_last_update_request()
99
100
101    def get_latest_initial_request(self):
102        """
103        Return the most recent initial update request.
104
105        AU requests occur in a chain of messages back and forth, e.g. the
106        initial request for an update -> the reply with the update -> the
107        report that install has started -> report that install has finished,
108        etc.  This function finds the first request in the latest such chain.
109
110        This message has no eventtype listed, or is rebooted_after_update
111        type (as an artifact from a previous update since this one).
112        Subsequent messages in the chain have different eventtype values.
113
114        @returns: string of the entire update request or None.
115
116        """
117        requests = self.get_update_requests()
118        if not requests:
119            return None
120
121        MATCH_STR = r'eventtype="(.*?)"'
122        for i in xrange(len(requests) - 1, -1, -1):
123            search = re.search(MATCH_STR, requests[i])
124            if (not search or
125                search.group(1) == _EVENT_TYPE_REBOOTED_AFTER_UPDATE):
126                return requests[i]
127
128        return None
129