1# Copyright (c) 2013 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 BaseHTTPServer
6import thread
7import urlparse
8
9from autotest_lib.client.bin import test, utils
10from autotest_lib.client.common_lib import error, utils
11
12def _split_url(url):
13    """Splits a URL into the URL base and path."""
14    split_url = urlparse.urlsplit(url)
15    url_base = urlparse.urlunsplit(
16            (split_url.scheme, split_url.netloc, '', '', ''))
17    url_path = split_url.path
18    return url_base, url_path.lstrip('/')
19
20class NanoOmahaDevserver(object):
21    """Simple implementation of Omaha."""
22
23    class Handler(BaseHTTPServer.BaseHTTPRequestHandler):
24        """Inner class for handling HTTP requests."""
25
26        _OMAHA_RESPONSE_TEMPLATE_HEAD = """
27          <response protocol=\"3.0\">
28            <daystart elapsed_seconds=\"44801\"/>
29            <app appid=\"{87efface-864d-49a5-9bb3-4b050a7c227a}\" status=\"ok\">
30              <ping status=\"ok\"/>
31              <updatecheck status=\"ok\">
32                <urls>
33                  <url codebase=\"%s\"/>
34                </urls>
35                <manifest version=\"9999.0.0\">
36                  <packages>
37                    <package name=\"%s\" size=\"%d\" required=\"true\"/>
38                  </packages>
39                  <actions>
40                    <action event=\"postinstall\"
41              ChromeOSVersion=\"9999.0.0\"
42              sha256=\"%s\"
43              needsadmin=\"false\"
44              IsDeltaPayload=\"false\"
45"""
46
47        _OMAHA_RESPONSE_TEMPLATE_TAIL = """ />
48                  </actions>
49                </manifest>
50              </updatecheck>
51            </app>
52          </response>
53"""
54
55        def do_POST(self):
56            """Handler for POST requests."""
57            if self.path == '/update':
58                (base, name) = _split_url(self.server._devserver._image_url)
59                response = self._OMAHA_RESPONSE_TEMPLATE_HEAD % (
60                        base + '/', name,
61                        self.server._devserver._image_size,
62                        self.server._devserver._sha256)
63                if self.server._devserver._metadata_size:
64                    response += '              MetadataSize="%d"\n' % (
65                            self.server._devserver._metadata_size)
66                if self.server._devserver._metadata_signature:
67                    response += '              MetadataSignatureRsa="%s"\n' % (
68                            self.server._devserver._metadata_signature)
69                if self.server._devserver._public_key:
70                    response += '              PublicKeyRsa="%s"\n' % (
71                            self.server._devserver._public_key)
72                response += self._OMAHA_RESPONSE_TEMPLATE_TAIL
73                self.send_response(200)
74                self.send_header('Content-Type', 'application/xml')
75                self.end_headers()
76                self.wfile.write(response)
77            else:
78                self.send_response(500)
79
80    def start(self):
81        """Starts the server."""
82        self._httpd = BaseHTTPServer.HTTPServer(('127.0.0.1', 0), self.Handler)
83        self._httpd._devserver = self
84        # Serve HTTP requests in a dedicated thread.
85        thread.start_new_thread(self._httpd.serve_forever, ())
86        self._port = self._httpd.socket.getsockname()[1]
87
88    def stop(self):
89        """Stops the server."""
90        self._httpd.shutdown()
91
92    def get_port(self):
93        """Returns the TCP port number the server is listening on."""
94        return self._port
95
96    def set_image_params(self, image_url, image_size, sha256,
97                         metadata_size=None,
98                         metadata_signature=None,
99                         public_key=None):
100        """Sets the values to return in the Omaha response. Only the
101        |image_url|, |image_size| and |sha256| parameters are
102        mandatory."""
103        self._image_url = image_url
104        self._image_size = image_size
105        self._sha256 = sha256
106        self._metadata_size = metadata_size
107        self._metadata_signature = metadata_signature
108        self._public_key = public_key
109
110
111class autoupdate_CannedOmahaUpdate(test.test):
112    """Client-side mechanism to update a DUT with a given image."""
113    version = 1
114
115    """Restarts update_engine and attempts an update from the image
116    pointed to by |image_url| of size |image_size| with checksum
117    |image_sha256|. The |metadata_size|, |metadata_signature| and
118    |public_key| parameters are optional.
119
120    If the |allow_failure| parameter is True, then the test will
121    succeed even if the update failed."""
122    def run_once(self, image_url, image_size, image_sha256,
123                 allow_failure=False,
124                 metadata_size=None,
125                 metadata_signature=None,
126                 public_key=None):
127        utils.run('restart update-engine')
128
129        omaha = NanoOmahaDevserver()
130        omaha.set_image_params(image_url,
131                               image_size,
132                               image_sha256,
133                               metadata_size,
134                               metadata_signature,
135                               public_key)
136        omaha.start()
137        try:
138            utils.run('update_engine_client -omaha_url=' +
139                      'http://127.0.0.1:%d/update ' % omaha.get_port() +
140                      '-update')
141        except error.CmdError:
142            omaha.stop()
143            if not allow_failure:
144                raise error.TestFail('Update attempt failed.')
145