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
5
6import base64
7import hashlib
8import httplib
9import json
10import logging
11import socket
12import StringIO
13import urllib2
14import urlparse
15
16try:
17    import pycurl
18except ImportError:
19    pycurl = None
20
21
22import common
23
24from autotest_lib.client.bin import utils
25from autotest_lib.client.common_lib import error
26from autotest_lib.client.common_lib.cros import retry
27from autotest_lib.server import frontend
28from autotest_lib.server import site_utils
29
30
31# Give all our rpcs about six seconds of retry time. If a longer timeout
32# is desired one should retry from the caller, this timeout is only meant
33# to avoid uncontrolled circumstances like network flake, not, say, retry
34# right across a reboot.
35BASE_REQUEST_TIMEOUT = 0.1
36JSON_HEADERS = {'Content-Type': 'application/json'}
37RPC_EXCEPTIONS = (httplib.BadStatusLine, socket.error, urllib2.HTTPError)
38MANIFEST_KEY = ('MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQC+hlN5FB+tjCsBszmBIvI'
39                'cD/djLLQm2zZfFygP4U4/o++ZM91EWtgII10LisoS47qT2TIOg4Un4+G57e'
40                'lZ9PjEIhcJfANqkYrD3t9dpEzMNr936TLB2u683B5qmbB68Nq1Eel7KVc+F'
41                '0BqhBondDqhvDvGPEV0vBsbErJFlNH7SQIDAQAB')
42SONIC_BOARD_LABEL = 'board:sonic'
43
44
45def get_extension_id(pub_key_pem=MANIFEST_KEY):
46    """Computes the extension id from the public key.
47
48    @param pub_key_pem: The public key used in the extension.
49
50    @return: The extension id.
51    """
52    pub_key_der = base64.b64decode(pub_key_pem)
53    sha = hashlib.sha256(pub_key_der).hexdigest()
54    prefix = sha[:32]
55    reencoded = ""
56    ord_a = ord('a')
57    for old_char in prefix:
58        code = int(old_char, 16)
59        new_char = chr(ord_a + code)
60        reencoded += new_char
61    return reencoded
62
63
64class Url(object):
65  """Container for URL information."""
66
67  def __init__(self):
68    self.scheme = 'http'
69    self.netloc = ''
70    self.path = ''
71    self.params = ''
72    self.query = ''
73    self.fragment = ''
74
75  def Build(self):
76    """Returns the URL."""
77    return urlparse.urlunparse((
78        self.scheme,
79        self.netloc,
80        self.path,
81        self.params,
82        self.query,
83        self.fragment))
84
85
86# TODO(beeps): Move get and post to curl too, since we have the need for
87# custom requests anyway.
88@retry.retry(RPC_EXCEPTIONS, timeout_min=BASE_REQUEST_TIMEOUT)
89def _curl_request(host, app_path, port, custom_request='', payload=None):
90    """Sends a custom request throug pycurl, to the url specified.
91    """
92    url = Url()
93    url.netloc = ':'.join((host, str(port)))
94    url.path = app_path
95    full_url = url.Build()
96
97    response = StringIO.StringIO()
98    conn = pycurl.Curl()
99    conn.setopt(conn.URL, full_url)
100    conn.setopt(conn.WRITEFUNCTION, response.write)
101    if custom_request:
102        conn.setopt(conn.CUSTOMREQUEST, custom_request)
103    if payload:
104        conn.setopt(conn.POSTFIELDS, payload)
105    conn.perform()
106    conn.close()
107    return response.getvalue()
108
109
110@retry.retry(RPC_EXCEPTIONS, timeout_min=BASE_REQUEST_TIMEOUT)
111def _get(url):
112    """Get request to the give url.
113
114    @raises: Any of the retry exceptions, if we hit the timeout.
115    @raises: error.TimeoutException, if the call itself times out.
116        eg: a hanging urlopen will get killed with a TimeoutException while
117        multiple retries that hit different Http errors will raise the last
118        HttpError instead of the TimeoutException.
119    """
120    return urllib2.urlopen(url).read()
121
122
123@retry.retry(RPC_EXCEPTIONS, timeout_min=BASE_REQUEST_TIMEOUT)
124def _post(url, data):
125    """Post data to the given url.
126
127    @param data: Json data to post.
128
129    @raises: Any of the retry exceptions, if we hit the timeout.
130    @raises: error.TimeoutException, if the call itself times out.
131        For examples see docstring for _get method.
132    """
133    request = urllib2.Request(url, json.dumps(data),
134                              headers=JSON_HEADERS)
135    urllib2.urlopen(request)
136
137
138@retry.retry(RPC_EXCEPTIONS + (error.TestError,), timeout_min=30)
139def acquire_sonic(lock_manager, additional_labels=None):
140    """Lock a host that has the sonic host labels.
141
142    @param lock_manager: A manager for locking/unlocking hosts, as defined by
143        server.cros.host_lock_manager.
144    @param additional_labels: A list of additional labels to apply in the search
145        for a sonic device.
146
147    @return: A string specifying the hostname of a locked sonic host.
148
149    @raises ValueError: Is no hosts matching the given labels are found.
150    """
151    sonic_host = None
152    afe = frontend.AFE(debug=True)
153    labels = [SONIC_BOARD_LABEL]
154    if additional_labels:
155        labels += additional_labels
156    sonic_hostname = utils.poll_for_condition(
157            lambda: site_utils.lock_host_with_labels(afe, lock_manager, labels),
158            sleep_interval=60,
159            exception=SonicProxyException('Timed out trying to find a sonic '
160                                          'host with labels %s.' % labels))
161    logging.info('Acquired sonic host returned %s', sonic_hostname)
162    return sonic_hostname
163
164
165class SonicProxyException(Exception):
166    """Generic exception raised when a sonic rpc fails."""
167    pass
168
169
170class SonicProxy(object):
171    """Client capable of making calls to the sonic device server."""
172    POLLING_INTERVAL = 5
173    SONIC_SERVER_PORT = '8008'
174
175    def __init__(self, hostname):
176        """
177        @param hostname: The name of the host for this sonic proxy.
178        """
179        self._sonic_server = 'http://%s:%s' % (hostname, self.SONIC_SERVER_PORT)
180        self._hostname = hostname
181
182
183    def check_server(self):
184        """Checks if the sonic server is up and running.
185
186        @raises: SonicProxyException if the server is unreachable.
187        """
188        try:
189            json.loads(_get(self._sonic_server))
190        except (RPC_EXCEPTIONS, error.TimeoutException) as e:
191            raise SonicProxyException('Could not retrieve information about '
192                                      'sonic device: %s' % e)
193
194
195    def reboot(self, when="now"):
196        """
197        Post to the server asking for a reboot.
198
199        @param when: The time till reboot. Can be any of:
200            now: immediately
201            fdr: set factory data reset flag and reboot now
202            ota: set recovery flag and reboot now
203            ota fdr: set both recovery and fdr flags, and reboot now
204            ota foreground: reboot and start force update page
205            idle: reboot only when idle screen usage > 10 mins
206
207        @raises SonicProxyException: if we're unable to post a reboot request.
208        """
209        reboot_url = '%s/%s/%s' % (self._sonic_server, 'setup', 'reboot')
210        reboot_params = {"params": when}
211        logging.info('Rebooting device through %s.', reboot_url)
212        try:
213            _post(reboot_url, reboot_params)
214        except (RPC_EXCEPTIONS, error.TimeoutException) as e:
215            raise SonicProxyException('Could not reboot sonic device through '
216                                      '%s: %s' % (self.SETUP_SERVER_PORT, e))
217
218
219    def stop_app(self, app):
220        """Stops the app.
221
222        Performs a hard reboot if pycurl isn't available.
223
224        @param app: An app name, eg YouTube, Fling, Netflix etc.
225
226        @raises pycurl.error: If the DELETE request fails after retries.
227        """
228        if not pycurl:
229            logging.warning('Rebooting sonic host to stop %s, please install '
230                            'pycurl if you do not wish to reboot.', app)
231            self.reboot()
232            return
233
234        _curl_request(self._hostname, 'apps/%s' % app,
235                      self.SONIC_SERVER_PORT, 'DELETE')
236
237
238    def start_app(self, app, payload):
239        """Starts an app.
240
241        @param app: An app name, eg YouTube, Fling, Netflix etc.
242        @param payload: An url payload for the app, eg: http://www.youtube.com.
243
244        @raises error.TimeoutException: If the call times out.
245        """
246        url = '%s/apps/%s' % (self._sonic_server, app)
247        _post(url, payload)
248
249