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 BaseHTTPServer
6import thread
7import urlparse
8
9def _split_url(url):
10    """Splits a URL into the URL base and path."""
11    split_url = urlparse.urlsplit(url)
12    url_base = urlparse.urlunsplit(
13            (split_url.scheme, split_url.netloc, '', '', ''))
14    url_path = split_url.path
15    return url_base, url_path.lstrip('/')
16
17
18class NanoOmahaDevserver(object):
19    """A simple Omaha instance that can be setup on a DUT in client tests."""
20
21    def __init__(self, eol=False):
22        """
23        Create a nano omaha devserver.
24
25        @param eol: True if we should return a response with _eol flag
26
27        """
28        self._eol = eol
29
30
31    class Handler(BaseHTTPServer.BaseHTTPRequestHandler):
32        """Inner class for handling HTTP requests."""
33
34        _OMAHA_RESPONSE_TEMPLATE_HEAD = """
35          <response protocol=\"3.0\">
36            <daystart elapsed_seconds=\"44801\"/>
37            <app appid=\"{87efface-864d-49a5-9bb3-4b050a7c227a}\" status=\"ok\">
38              <ping status=\"ok\"/>
39              <updatecheck status=\"ok\">
40                <urls>
41                  <url codebase=\"%s\"/>
42                </urls>
43                <manifest version=\"9999.0.0\">
44                  <packages>
45                    <package name=\"%s\" size=\"%d\" required=\"true\"/>
46                  </packages>
47                  <actions>
48                    <action event=\"postinstall\"
49              ChromeOSVersion=\"9999.0.0\"
50              sha256=\"%s\"
51              needsadmin=\"false\"
52              IsDeltaPayload=\"%s\"
53        """
54
55        _OMAHA_RESPONSE_TEMPLATE_TAIL = """ />
56                  </actions>
57                </manifest>
58              </updatecheck>
59            </app>
60          </response>
61        """
62
63        _OMAHA_RESPONSE_EOL = """
64          <response protocol=\"3.0\">
65            <daystart elapsed_seconds=\"44801\"/>
66            <app appid=\"{87efface-864d-49a5-9bb3-4b050a7c227a}\" status=\"ok\">
67              <ping status=\"ok\"/>
68              <updatecheck _eol=\"eol\" status=\"noupdate\"/>
69            </app>
70          </response>
71        """
72
73        def do_POST(self):
74            """Handler for POST requests."""
75            if self.path == '/update':
76                if self.server._devserver._eol:
77                    response = self._OMAHA_RESPONSE_EOL
78                else:
79                    (base, name) = _split_url(self.server._devserver._image_url)
80                    response = self._OMAHA_RESPONSE_TEMPLATE_HEAD % (
81                            base + '/',
82                            name,
83                            self.server._devserver._image_size,
84                            self.server._devserver._sha256,
85                            str(self.server._devserver._is_delta).lower())
86                    if self.server._devserver._is_delta:
87                        response += '              IsDelta="true"\n'
88                    if self.server._devserver._critical:
89                        response += '              deadline="now"\n'
90                    if self.server._devserver._metadata_size:
91                        response += '              MetadataSize="%d"\n' % (
92                                self.server._devserver._metadata_size)
93                    if self.server._devserver._metadata_signature:
94                        response += '              ' \
95                                    'MetadataSignatureRsa="%s"\n' % (
96                                self.server._devserver._metadata_signature)
97                    if self.server._devserver._public_key:
98                        response += '              PublicKeyRsa="%s"\n' % (
99                                self.server._devserver._public_key)
100                    response += self._OMAHA_RESPONSE_TEMPLATE_TAIL
101                self.send_response(200)
102                self.send_header('Content-Type', 'application/xml')
103                self.end_headers()
104                self.wfile.write(response)
105            else:
106                self.send_response(500)
107
108    def start(self):
109        """Starts the server."""
110        self._httpd = BaseHTTPServer.HTTPServer(('127.0.0.1', 0), self.Handler)
111        self._httpd._devserver = self
112        # Serve HTTP requests in a dedicated thread.
113        thread.start_new_thread(self._httpd.serve_forever, ())
114        self._port = self._httpd.socket.getsockname()[1]
115
116    def stop(self):
117        """Stops the server."""
118        self._httpd.shutdown()
119
120    def get_port(self):
121        """Returns the TCP port number the server is listening on."""
122        return self._port
123
124    def set_image_params(self, image_url, image_size, sha256,
125                         metadata_size=None, metadata_signature=None,
126                         public_key=None, is_delta=False, critical=False):
127        """Sets the values to return in the Omaha response. Only the
128        |image_url|, |image_size| and |sha256| parameters are
129        mandatory."""
130        self._image_url = image_url
131        self._image_size = image_size
132        self._sha256 = sha256
133        self._metadata_size = metadata_size
134        self._metadata_signature = metadata_signature
135        self._public_key = public_key
136        self._is_delta = is_delta
137        self._critical = critical
138