• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright (c) 2016 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"""
6This module includes all moblab-related RPCs. These RPCs can only be run
7on moblab.
8"""
9
10import ConfigParser
11import common
12import logging
13import os
14import re
15import shutil
16import socket
17import StringIO
18import subprocess
19
20from autotest_lib.client.common_lib import error
21from autotest_lib.client.common_lib import global_config
22from autotest_lib.client.common_lib import utils
23from autotest_lib.frontend.afe import models
24from autotest_lib.frontend.afe import rpc_utils
25from autotest_lib.server import frontend
26from autotest_lib.server.hosts import moblab_host
27from chromite.lib import gs
28
29_CONFIG = global_config.global_config
30MOBLAB_BOTO_LOCATION = '/home/moblab/.boto'
31CROS_CACHEDIR = '/mnt/moblab/cros_cache_apache'
32
33# Google Cloud Storage bucket url regex pattern. The pattern is used to extract
34# the bucket name from the bucket URL. For example, "gs://image_bucket/google"
35# should result in a bucket name "image_bucket".
36GOOGLE_STORAGE_BUCKET_URL_PATTERN = re.compile(
37        r'gs://(?P<bucket>[a-zA-Z][a-zA-Z0-9-_]*)/?.*')
38
39# Contants used in Json RPC field names.
40_IMAGE_STORAGE_SERVER = 'image_storage_server'
41_GS_ACCESS_KEY_ID = 'gs_access_key_id'
42_GS_SECRET_ACCESS_KEY = 'gs_secret_access_key'
43_RESULT_STORAGE_SERVER = 'results_storage_server'
44_USE_EXISTING_BOTO_FILE = 'use_existing_boto_file'
45_CLOUD_NOTIFICATION_ENABLED = 'cloud_notification_enabled'
46
47# Location where dhcp leases are stored.
48_DHCPD_LEASES = '/var/lib/dhcp/dhcpd.leases'
49
50# File where information about the current device is stored.
51_ETC_LSB_RELEASE = '/etc/lsb-release'
52
53# Full path to the correct gsutil command to run.
54class GsUtil:
55    _GSUTIL_CMD = None
56
57    @classmethod
58    def get_gsutil_cmd(cls):
59      if not cls._GSUTIL_CMD:
60         cls._GSUTIL_CMD = gs.GSContext.GetDefaultGSUtilBin(
61           cache_dir=CROS_CACHEDIR)
62
63      return cls._GSUTIL_CMD
64
65
66class BucketPerformanceTestException(Exception):
67  pass
68
69@rpc_utils.moblab_only
70def get_config_values():
71    """Returns all config values parsed from global and shadow configs.
72
73    Config values are grouped by sections, and each section is composed of
74    a list of name value pairs.
75    """
76    sections =_CONFIG.get_sections()
77    config_values = {}
78    for section in sections:
79        config_values[section] = _CONFIG.config.items(section)
80    return rpc_utils.prepare_for_serialization(config_values)
81
82
83def _write_config_file(config_file, config_values, overwrite=False):
84    """Writes out a configuration file.
85
86    @param config_file: The name of the configuration file.
87    @param config_values: The ConfigParser object.
88    @param ovewrite: Flag on if overwriting is allowed.
89    """
90    if not config_file:
91        raise error.RPCException('Empty config file name.')
92    if not overwrite and os.path.exists(config_file):
93        raise error.RPCException('Config file already exists.')
94
95    if config_values:
96        with open(config_file, 'w') as config_file:
97            config_values.write(config_file)
98
99
100def _read_original_config():
101    """Reads the orginal configuratino without shadow.
102
103    @return: A configuration object, see global_config_class.
104    """
105    original_config = global_config.global_config_class()
106    original_config.set_config_files(shadow_file='')
107    return original_config
108
109
110def _read_raw_config(config_file):
111    """Reads the raw configuration from a configuration file.
112
113    @param: config_file: The path of the configuration file.
114
115    @return: A ConfigParser object.
116    """
117    shadow_config = ConfigParser.RawConfigParser()
118    shadow_config.read(config_file)
119    return shadow_config
120
121
122def _get_shadow_config_from_partial_update(config_values):
123    """Finds out the new shadow configuration based on a partial update.
124
125    Since the input is only a partial config, we should not lose the config
126    data inside the existing shadow config file. We also need to distinguish
127    if the input config info overrides with a new value or reverts back to
128    an original value.
129
130    @param config_values: See get_moblab_settings().
131
132    @return: The new shadow configuration as ConfigParser object.
133    """
134    original_config = _read_original_config()
135    existing_shadow = _read_raw_config(_CONFIG.shadow_file)
136    for section, config_value_list in config_values.iteritems():
137        for key, value in config_value_list:
138            if original_config.get_config_value(section, key,
139                                                default='',
140                                                allow_blank=True) != value:
141                if not existing_shadow.has_section(section):
142                    existing_shadow.add_section(section)
143                existing_shadow.set(section, key, value)
144            elif existing_shadow.has_option(section, key):
145                existing_shadow.remove_option(section, key)
146    return existing_shadow
147
148
149def _update_partial_config(config_values):
150    """Updates the shadow configuration file with a partial config udpate.
151
152    @param config_values: See get_moblab_settings().
153    """
154    existing_config = _get_shadow_config_from_partial_update(config_values)
155    _write_config_file(_CONFIG.shadow_file, existing_config, True)
156
157
158@rpc_utils.moblab_only
159def update_config_handler(config_values):
160    """Update config values and override shadow config.
161
162    @param config_values: See get_moblab_settings().
163    """
164    original_config = _read_original_config()
165    new_shadow = ConfigParser.RawConfigParser()
166    for section, config_value_list in config_values.iteritems():
167        for key, value in config_value_list:
168            if original_config.get_config_value(section, key,
169                                                default='',
170                                                allow_blank=True) != value:
171                if not new_shadow.has_section(section):
172                    new_shadow.add_section(section)
173                new_shadow.set(section, key, value)
174
175    if not _CONFIG.shadow_file or not os.path.exists(_CONFIG.shadow_file):
176        raise error.RPCException('Shadow config file does not exist.')
177    _write_config_file(_CONFIG.shadow_file, new_shadow, True)
178
179    # TODO (sbasi) crbug.com/403916 - Remove the reboot command and
180    # instead restart the services that rely on the config values.
181    os.system('sudo reboot')
182
183
184@rpc_utils.moblab_only
185def reset_config_settings():
186    """Reset moblab shadow config."""
187    with open(_CONFIG.shadow_file, 'w') as config_file:
188        pass
189    os.system('sudo reboot')
190
191
192@rpc_utils.moblab_only
193def reboot_moblab():
194    """Simply reboot the device."""
195    os.system('sudo reboot')
196
197
198@rpc_utils.moblab_only
199def set_boto_key(boto_key):
200    """Update the boto_key file.
201
202    @param boto_key: File name of boto_key uploaded through handle_file_upload.
203    """
204    if not os.path.exists(boto_key):
205        raise error.RPCException('Boto key: %s does not exist!' % boto_key)
206    shutil.copyfile(boto_key, moblab_host.MOBLAB_BOTO_LOCATION)
207
208
209@rpc_utils.moblab_only
210def set_service_account_credential(service_account_filename):
211    """Update the service account credential file.
212
213    @param service_account_filename: Name of uploaded file through
214            handle_file_upload.
215    """
216    if not os.path.exists(service_account_filename):
217        raise error.RPCException(
218                'Service account file: %s does not exist!' %
219                service_account_filename)
220    shutil.copyfile(
221            service_account_filename,
222            moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION)
223
224
225@rpc_utils.moblab_only
226def set_launch_control_key(launch_control_key):
227    """Update the launch_control_key file.
228
229    @param launch_control_key: File name of launch_control_key uploaded through
230            handle_file_upload.
231    """
232    if not os.path.exists(launch_control_key):
233        raise error.RPCException('Launch Control key: %s does not exist!' %
234                                 launch_control_key)
235    shutil.copyfile(launch_control_key,
236                    moblab_host.MOBLAB_LAUNCH_CONTROL_KEY_LOCATION)
237    # Restart the devserver service.
238    os.system('sudo restart moblab-devserver-init')
239
240
241###########Moblab Config Wizard RPCs #######################
242def _get_public_ip_address(socket_handle):
243    """Gets the public IP address.
244
245    Connects to Google DNS server using a socket and gets the preferred IP
246    address from the connection.
247
248    @param: socket_handle: a unix socket.
249
250    @return: public ip address as string.
251    """
252    try:
253        socket_handle.settimeout(1)
254        socket_handle.connect(('8.8.8.8', 53))
255        socket_name = socket_handle.getsockname()
256        if socket_name is not None:
257            logging.info('Got socket name from UDP socket.')
258            return socket_name[0]
259        logging.warn('Created UDP socket but with no socket_name.')
260    except socket.error:
261        logging.warn('Could not get socket name from UDP socket.')
262    return None
263
264
265def _get_network_info():
266    """Gets the network information.
267
268    TCP socket is used to test the connectivity. If there is no connectivity,
269    try to get the public IP with UDP socket.
270
271    @return: a tuple as (public_ip_address, connected_to_internet).
272    """
273    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
274    ip = _get_public_ip_address(s)
275    if ip is not None:
276        logging.info('Established TCP connection with well known server.')
277        return (ip, True)
278    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
279    return (_get_public_ip_address(s), False)
280
281
282@rpc_utils.moblab_only
283def get_network_info():
284    """Returns the server ip addresses, and if the server connectivity.
285
286    The server ip addresses as an array of strings, and the connectivity as a
287    flag.
288    """
289    network_info = {}
290    info = _get_network_info()
291    if info[0] is not None:
292        network_info['server_ips'] = [info[0]]
293    network_info['is_connected'] = info[1]
294
295    return rpc_utils.prepare_for_serialization(network_info)
296
297
298# Gets the boto configuration.
299def _get_boto_config():
300    """Reads the boto configuration from the boto file.
301
302    @return: Boto configuration as ConfigParser object.
303    """
304    boto_config = ConfigParser.ConfigParser()
305    boto_config.read(MOBLAB_BOTO_LOCATION)
306    return boto_config
307
308
309@rpc_utils.moblab_only
310def get_cloud_storage_info():
311    """RPC handler to get the cloud storage access information.
312    """
313    cloud_storage_info = {}
314    value =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
315    if value is not None:
316        cloud_storage_info[_IMAGE_STORAGE_SERVER] = value
317    value = _CONFIG.get_config_value('CROS', _RESULT_STORAGE_SERVER,
318            default=None)
319    if value is not None:
320        cloud_storage_info[_RESULT_STORAGE_SERVER] = value
321
322    boto_config = _get_boto_config()
323    sections = boto_config.sections()
324
325    if sections:
326        cloud_storage_info[_USE_EXISTING_BOTO_FILE] = True
327    else:
328        cloud_storage_info[_USE_EXISTING_BOTO_FILE] = False
329    if 'Credentials' in sections:
330        options = boto_config.options('Credentials')
331        if _GS_ACCESS_KEY_ID in options:
332            value = boto_config.get('Credentials', _GS_ACCESS_KEY_ID)
333            cloud_storage_info[_GS_ACCESS_KEY_ID] = value
334        if _GS_SECRET_ACCESS_KEY in options:
335            value = boto_config.get('Credentials', _GS_SECRET_ACCESS_KEY)
336            cloud_storage_info[_GS_SECRET_ACCESS_KEY] = value
337
338    return rpc_utils.prepare_for_serialization(cloud_storage_info)
339
340
341def _get_bucket_name_from_url(bucket_url):
342    """Gets the bucket name from a bucket url.
343
344    @param: bucket_url: the bucket url string.
345    """
346    if bucket_url:
347        match = GOOGLE_STORAGE_BUCKET_URL_PATTERN.match(bucket_url)
348        if match:
349            return match.group('bucket')
350    return None
351
352def _is_valid_boto_key(key_id, key_secret, directory):
353  try:
354      _run_bucket_performance_test(key_id, key_secret, directory)
355  except BucketPerformanceTestException as e:
356       return(False, str(e))
357  return(True, None)
358
359def _validate_cloud_storage_info(cloud_storage_info):
360    """Checks if the cloud storage information is valid.
361
362    @param: cloud_storage_info: The JSON RPC object for cloud storage info.
363
364    @return: A tuple as (valid_boolean, details_string).
365    """
366    valid = True
367    details = None
368    if not cloud_storage_info[_USE_EXISTING_BOTO_FILE]:
369        key_id = cloud_storage_info[_GS_ACCESS_KEY_ID]
370        key_secret = cloud_storage_info[_GS_SECRET_ACCESS_KEY]
371        valid, details = _is_valid_boto_key(
372            key_id, key_secret, cloud_storage_info[_IMAGE_STORAGE_SERVER])
373    return (valid, details)
374
375
376def _create_operation_status_response(is_ok, details):
377    """Helper method to create a operation status reponse.
378
379    @param: is_ok: Boolean for if the operation is ok.
380    @param: details: A detailed string.
381
382    @return: A serialized JSON RPC object.
383    """
384    status_response = {'status_ok': is_ok}
385    if details:
386        status_response['status_details'] = details
387    return rpc_utils.prepare_for_serialization(status_response)
388
389
390@rpc_utils.moblab_only
391def validate_cloud_storage_info(cloud_storage_info):
392    """RPC handler to check if the cloud storage info is valid.
393
394    @param cloud_storage_info: The JSON RPC object for cloud storage info.
395    """
396    valid, details = _validate_cloud_storage_info(cloud_storage_info)
397    return _create_operation_status_response(valid, details)
398
399
400@rpc_utils.moblab_only
401def submit_wizard_config_info(cloud_storage_info):
402    """RPC handler to submit the cloud storage info.
403
404    @param cloud_storage_info: The JSON RPC object for cloud storage info.
405    """
406    config_update = {}
407    config_update['CROS'] = [
408        (_IMAGE_STORAGE_SERVER, cloud_storage_info[_IMAGE_STORAGE_SERVER]),
409        (_RESULT_STORAGE_SERVER, cloud_storage_info[_RESULT_STORAGE_SERVER])
410    ]
411    _update_partial_config(config_update)
412
413    if not cloud_storage_info[_USE_EXISTING_BOTO_FILE]:
414        boto_config = ConfigParser.RawConfigParser()
415        boto_config.add_section('Credentials')
416        boto_config.set('Credentials', _GS_ACCESS_KEY_ID,
417                        cloud_storage_info[_GS_ACCESS_KEY_ID])
418        boto_config.set('Credentials', _GS_SECRET_ACCESS_KEY,
419                        cloud_storage_info[_GS_SECRET_ACCESS_KEY])
420        _write_config_file(MOBLAB_BOTO_LOCATION, boto_config, True)
421
422    _CONFIG.parse_config_file()
423    _enable_notification_using_credentials_in_bucket()
424    services = ['moblab-devserver-init',
425    'moblab-devserver-cleanup-init', 'moblab-gsoffloader_s-init',
426    'moblab-scheduler-init', 'moblab-gsoffloader-init']
427    cmd = 'export ATEST_RESULTS_DIR=/usr/local/autotest/results;'
428    cmd += 'sudo stop ' + ';sudo stop '.join(services)
429    cmd += ';sudo start ' + ';sudo start '.join(services)
430    cmd += ';sudo apache2 -k graceful'
431    logging.info(cmd)
432    try:
433        utils.run(cmd)
434    except error.CmdError as e:
435        logging.error(e)
436        # if all else fails reboot the device.
437        utils.run('sudo reboot')
438
439    return _create_operation_status_response(True, None)
440
441
442@rpc_utils.moblab_only
443def get_version_info():
444    """ RPC handler to get informaiton about the version of the moblab.
445
446    @return: A serialized JSON RPC object.
447    """
448    lines = open(_ETC_LSB_RELEASE).readlines()
449    version_response = {
450        x.split('=')[0]: x.split('=')[1] for x in lines if '=' in x}
451    version_response['MOBLAB_ID'] = utils.get_moblab_id();
452    version_response['MOBLAB_MAC_ADDRESS'] = (
453        utils.get_default_interface_mac_address())
454    return rpc_utils.prepare_for_serialization(version_response)
455
456
457@rpc_utils.moblab_only
458def get_connected_dut_info():
459    """ RPC handler to get informaiton about the DUTs connected to the moblab.
460
461    @return: A serialized JSON RPC object.
462    """
463    # Make a list of the connected DUT's
464    leases = _get_dhcp_dut_leases()
465
466    # Get a list of the AFE configured DUT's
467    hosts = list(rpc_utils.get_host_query((), False, True, {}))
468    models.Host.objects.populate_relationships(hosts, models.Label,
469                                               'label_list')
470    configured_duts = {}
471    for host in hosts:
472        labels = [label.name for label in host.label_list]
473        labels.sort()
474        configured_duts[host.hostname] = ', '.join(labels)
475
476    return rpc_utils.prepare_for_serialization(
477            {'configured_duts': configured_duts,
478             'connected_duts': leases})
479
480
481def _get_dhcp_dut_leases():
482     """ Extract information about connected duts from the dhcp server.
483
484     @return: A dict of ipaddress to mac address for each device connected.
485     """
486     lease_info = open(_DHCPD_LEASES).read()
487
488     leases = {}
489     for lease in lease_info.split('lease'):
490         if lease.find('binding state active;') != -1:
491             ipaddress = lease.split('\n')[0].strip(' {')
492             last_octet = int(ipaddress.split('.')[-1].strip())
493             if last_octet > 150:
494                 continue
495             mac_address_search = re.search('hardware ethernet (.*);', lease)
496             if mac_address_search:
497                 leases[ipaddress] = mac_address_search.group(1)
498     return leases
499
500
501@rpc_utils.moblab_only
502def add_moblab_dut(ipaddress):
503    """ RPC handler to add a connected DUT to autotest.
504
505    @param ipaddress: IP address of the DUT.
506
507    @return: A string giving information about the status.
508    """
509    cmd = '/usr/local/autotest/cli/atest host create %s &' % ipaddress
510    subprocess.call(cmd, shell=True)
511    return (True, 'DUT %s added to Autotest' % ipaddress)
512
513
514@rpc_utils.moblab_only
515def remove_moblab_dut(ipaddress):
516    """ RPC handler to remove DUT entry from autotest.
517
518    @param ipaddress: IP address of the DUT.
519
520    @return: True if the command succeeds without an exception
521    """
522    models.Host.smart_get(ipaddress).delete()
523    return (True, 'DUT %s deleted from Autotest' % ipaddress)
524
525
526@rpc_utils.moblab_only
527def add_moblab_label(ipaddress, label_name):
528    """ RPC handler to add a label in autotest to a DUT entry.
529
530    @param ipaddress: IP address of the DUT.
531    @param label_name: The label name.
532
533    @return: A string giving information about the status.
534    """
535    # Try to create the label in case it does not already exist.
536    label = None
537    try:
538        label = models.Label.add_object(name=label_name)
539    except:
540        label = models.Label.smart_get(label_name)
541    host_obj = models.Host.smart_get(ipaddress)
542    if label:
543        label.host_set.add(host_obj)
544        return (True, 'Added label %s to DUT %s' % (label_name, ipaddress))
545    return (False,
546            'Failed to add label %s to DUT %s' % (label_name, ipaddress))
547
548
549@rpc_utils.moblab_only
550def remove_moblab_label(ipaddress, label_name):
551    """ RPC handler to remove a label in autotest from a DUT entry.
552
553    @param ipaddress: IP address of the DUT.
554    @param label_name: The label name.
555
556    @return: A string giving information about the status.
557    """
558    host_obj = models.Host.smart_get(ipaddress)
559    models.Label.smart_get(label_name).host_set.remove(host_obj)
560    return (True, 'Removed label %s from DUT %s' % (label_name, ipaddress))
561
562
563def _get_connected_dut_labels(requested_label, only_first_label=True):
564    """ Query the DUT's attached to the moblab and return a filtered list
565        of labels.
566
567    @param requested_label:  the label name you are requesting.
568    @param only_first_label:  if the device has the same label name multiple
569                              times only return the first label value in the
570                              list.
571
572    @return: A de-duped list of requested dut labels attached to the moblab.
573    """
574    hosts = list(rpc_utils.get_host_query((), False, True, {}))
575    if not hosts:
576        return []
577    models.Host.objects.populate_relationships(hosts, models.Label,
578                                               'label_list')
579    labels = set()
580    for host in hosts:
581        for label in host.label_list:
582            if requested_label in label.name:
583                labels.add(label.name.replace(requested_label, ''))
584                if only_first_label:
585                    break
586    return list(labels)
587
588
589@rpc_utils.moblab_only
590def get_connected_boards():
591    """ RPC handler to get a list of the boards connected to the moblab.
592
593    @return: A de-duped list of board types attached to the moblab.
594    """
595    boards = _get_connected_dut_labels("board:")
596    boards.sort()
597    return boards
598
599
600@rpc_utils.moblab_only
601def get_connected_pools():
602    """ RPC handler to get a list of the pools labels on the DUT's connected.
603
604    @return: A de-duped list of pool labels.
605    """
606    pools = _get_connected_dut_labels("pool:", False)
607    pools.sort()
608    return pools
609
610
611@rpc_utils.moblab_only
612def get_builds_for_board(board_name):
613    """ RPC handler to find the most recent builds for a board.
614
615
616    @param board_name: The name of a connected board.
617    @return: A list of string with the most recent builds for the latest
618             three milestones.
619    """
620    return _get_builds_for_in_directory(board_name + '-release')
621
622
623@rpc_utils.moblab_only
624def get_firmware_for_board(board_name):
625    """ RPC handler to find the most recent firmware for a board.
626
627
628    @param board_name: The name of a connected board.
629    @return: A list of strings with the most recent firmware builds for the
630             latest three milestones.
631    """
632    return _get_builds_for_in_directory(board_name + '-firmware')
633
634
635def _get_sortable_build_number(sort_key):
636    """ Converts a build number line cyan-release/R59-9460.27.0 into an integer.
637
638        To be able to sort a list of builds you need to convert the build number
639        into an integer so it can be compared correctly to other build.
640
641        cyan-release/R59-9460.27.0 =>  5909460027000
642
643        If the sort key is not recognised as a build number 1 will be returned.
644
645    @param sort_key: A string that represents a build number like
646                     cyan-release/R59-9460.27.0
647    @return: An integer that represents that build number or 1 if not recognised
648             as a build.
649    """
650    build_number = re.search('.*/R([0-9]*)-([0-9]*)\.([0-9]*)\.([0-9]*)',
651                             sort_key)
652    if not build_number or not len(build_number.groups()) == 4:
653      return 1
654    return int("%d%05d%03d%03d" % (int(build_number.group(1)),
655                                   int(build_number.group(2)),
656                                   int(build_number.group(3)),
657                                   int(build_number.group(4))))
658
659def _get_builds_for_in_directory(directory_name, milestone_limit=3,
660                                 build_limit=20):
661    """ Fetch the most recent builds for the last three milestones from gcs.
662
663
664    @param directory_name: The sub-directory under the configured GCS image
665                           storage bucket to search.
666
667
668    @return: A string list no longer than <milestone_limit> x <build_limit>
669             items, containing the most recent <build_limit> builds from the
670             last milestone_limit milestones.
671    """
672    output = StringIO.StringIO()
673    gs_image_location =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
674    utils.run(GsUtil.get_gsutil_cmd(),
675              args=('ls', gs_image_location + directory_name),
676              stdout_tee=output)
677    lines = output.getvalue().split('\n')
678    output.close()
679    builds = [line.replace(gs_image_location,'').strip('/ ')
680              for line in lines if line != '']
681    build_matcher = re.compile(r'^.*\/R([0-9]*)-.*')
682    build_map = {}
683    for build in builds:
684        match = build_matcher.match(build)
685        if match:
686            milestone = match.group(1)
687            if milestone not in build_map:
688                build_map[milestone] = []
689            build_map[milestone].append(build)
690    milestones = build_map.keys()
691    milestones.sort()
692    milestones.reverse()
693    build_list = []
694    for milestone in milestones[:milestone_limit]:
695         builds = build_map[milestone]
696         builds.sort(key=_get_sortable_build_number)
697         builds.reverse()
698         build_list.extend(builds[:build_limit])
699    return build_list
700
701
702def _run_bucket_performance_test(key_id, key_secret, bucket_name,
703                                 test_size='1M', iterations='1',
704                                 result_file='/tmp/gsutil_perf.json'):
705    """Run a gsutil perfdiag on a supplied bucket and output the results"
706
707       @param key_id: boto key of the bucket to be accessed
708       @param key_secret: boto secret of the bucket to be accessed
709       @param bucket_name: bucket to be tested.
710       @param test_size: size of file to use in test, see gsutil perfdiag help.
711       @param iterations: number of times each test is run.
712       @param result_file: name of file to write results out to.
713
714       @return None
715       @raises BucketPerformanceTestException if the command fails.
716    """
717    try:
718      utils.run(GsUtil.get_gsutil_cmd(), args=(
719          '-o', 'Credentials:gs_access_key_id=%s' % key_id,
720          '-o', 'Credentials:gs_secret_access_key=%s' % key_secret,
721          'perfdiag', '-s', test_size, '-o', result_file,
722          '-n', iterations,
723          bucket_name))
724    except error.CmdError as e:
725       logging.error(e)
726       # Extract useful error from the stacktrace
727       errormsg = str(e)
728       start_error_pos = errormsg.find("<Error>")
729       end_error_pos = errormsg.find("</Error>", start_error_pos)
730       extracted_error_msg = errormsg[start_error_pos:end_error_pos]
731       raise BucketPerformanceTestException(
732           extracted_error_msg if extracted_error_msg else errormsg)
733    # TODO(haddowk) send the results to the cloud console when that feature is
734    # enabled.
735
736
737@rpc_utils.moblab_only
738def run_suite(board, build, suite, ro_firmware=None, rw_firmware=None,
739              pool=None, suite_args=None):
740    """ RPC handler to run a test suite.
741
742    @param board: a board name connected to the moblab.
743    @param build: a build name of a build in the GCS.
744    @param suite: the name of a suite to run
745    @param ro_firmware: Optional ro firmware build number to use.
746    @param rw_firmware: Optional rw firmware build number to use.
747    @param pool: Optional pool name to run the suite in.
748    @param suite_args: Arguments to be used in the suite control file.
749
750    @return: None
751    """
752    builds = {'cros-version': build}
753    if rw_firmware:
754        builds['fwrw-version'] = rw_firmware
755    if ro_firmware:
756        builds['fwro-version'] = ro_firmware
757    if suite_args:
758        list_suite_args = map(lambda s: s.strip(), suite_args.split(','))
759    else:
760        list_suite_args = None
761    afe = frontend.AFE(user='moblab')
762    afe.run('create_suite_job', board=board, builds=builds, name=suite,
763    pool=pool, run_prod_code=False, test_source_build=build,
764    wait_for_results=False, suite_args=list_suite_args)
765
766
767def _enable_notification_using_credentials_in_bucket():
768    """ Check and enable cloud notification if a credentials file exits.
769    @return: None
770    """
771    gs_image_location =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
772    try:
773        utils.run(GsUtil.get_gsutil_cmd(), args=(
774            'cp', gs_image_location + 'pubsub-key-do-not-delete.json', '/tmp'))
775        # This runs the copy as moblab user
776        shutil.copyfile('/tmp/pubsub-key-do-not-delete.json',
777                        moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION)
778
779    except error.CmdError as e:
780        logging.error(e)
781    else:
782        logging.info('Enabling cloud notifications')
783        config_update = {}
784        config_update['CROS'] = [(_CLOUD_NOTIFICATION_ENABLED, True)]
785        _update_partial_config(config_update)
786