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 sys
16import shutil
17import socket
18import StringIO
19import subprocess
20import time
21import multiprocessing
22import ctypes
23
24from autotest_lib.client.common_lib import error
25from autotest_lib.client.common_lib import global_config
26from autotest_lib.client.common_lib import utils
27from autotest_lib.frontend.afe import models
28from autotest_lib.frontend.afe import rpc_utils
29from autotest_lib.server import frontend
30from autotest_lib.server.hosts import moblab_host
31
32_CONFIG = global_config.global_config
33MOBLAB_BOTO_LOCATION = '/home/moblab/.boto'
34CROS_CACHEDIR = '/mnt/moblab/cros_cache_apache'
35
36# Google Cloud Storage bucket url regex pattern. The pattern is used to extract
37# the bucket name from the bucket URL. For example, "gs://image_bucket/google"
38# should result in a bucket name "image_bucket".
39GOOGLE_STORAGE_BUCKET_URL_PATTERN = re.compile(
40        r'gs://(?P<bucket>[a-zA-Z][a-zA-Z0-9-_]*)/?.*')
41
42# Contants used in Json RPC field names.
43_IMAGE_STORAGE_SERVER = 'image_storage_server'
44_GS_ACCESS_KEY_ID = 'gs_access_key_id'
45_GS_SECRET_ACCESS_KEY = 'gs_secret_access_key'
46_RESULT_STORAGE_SERVER = 'results_storage_server'
47_USE_EXISTING_BOTO_FILE = 'use_existing_boto_file'
48_CLOUD_NOTIFICATION_ENABLED = 'cloud_notification_enabled'
49_WIFI_AP_NAME = 'wifi_dut_ap_name'
50_WIFI_AP_PASS = 'wifi_dut_ap_pass'
51
52# Location where dhcp leases are stored.
53_DHCPD_LEASES = '/var/lib/dhcp/dhcpd.leases'
54
55# File where information about the current device is stored.
56_ETC_LSB_RELEASE = '/etc/lsb-release'
57
58# ChromeOS update engine client binary location
59_UPDATE_ENGINE_CLIENT = '/usr/bin/update_engine_client'
60
61# Set the suite timeout per suite in minutes
62# default is 24 hours
63_DEFAULT_SUITE_TIMEOUT_MINS = 1440
64_SUITE_TIMEOUT_MAP = {
65    'hardware_storagequal': 40320,
66    'hardware_storagequal_quick': 40320
67}
68
69# Full path to the correct gsutil command to run.
70class GsUtil:
71    """Helper class to find correct gsutil command."""
72    _GSUTIL_CMD = None
73
74    @classmethod
75    def get_gsutil_cmd(cls):
76      if not cls._GSUTIL_CMD:
77         cls._GSUTIL_CMD = 'gsutil'
78
79      return cls._GSUTIL_CMD
80
81
82class BucketPerformanceTestException(Exception):
83  """Exception thrown when the command to test the bucket performance fails."""
84  pass
85
86@rpc_utils.moblab_only
87def get_config_values():
88    """Returns all config values parsed from global and shadow configs.
89
90    Config values are grouped by sections, and each section is composed of
91    a list of name value pairs.
92    """
93    sections =_CONFIG.get_sections()
94    config_values = {}
95    for section in sections:
96        config_values[section] = _CONFIG.config.items(section)
97    return rpc_utils.prepare_for_serialization(config_values)
98
99
100def _write_config_file(config_file, config_values, overwrite=False):
101    """Writes out a configuration file.
102
103    @param config_file: The name of the configuration file.
104    @param config_values: The ConfigParser object.
105    @param ovewrite: Flag on if overwriting is allowed.
106    """
107    if not config_file:
108        raise error.RPCException('Empty config file name.')
109    if not overwrite and os.path.exists(config_file):
110        raise error.RPCException('Config file already exists.')
111
112    if config_values:
113        with open(config_file, 'w') as config_file:
114            config_values.write(config_file)
115
116
117def _read_original_config():
118    """Reads the orginal configuratino without shadow.
119
120    @return: A configuration object, see global_config_class.
121    """
122    original_config = global_config.global_config_class()
123    original_config.set_config_files(shadow_file='')
124    return original_config
125
126
127def _read_raw_config(config_file):
128    """Reads the raw configuration from a configuration file.
129
130    @param: config_file: The path of the configuration file.
131
132    @return: A ConfigParser object.
133    """
134    shadow_config = ConfigParser.RawConfigParser()
135    shadow_config.read(config_file)
136    return shadow_config
137
138
139def _get_shadow_config_from_partial_update(config_values):
140    """Finds out the new shadow configuration based on a partial update.
141
142    Since the input is only a partial config, we should not lose the config
143    data inside the existing shadow config file. We also need to distinguish
144    if the input config info overrides with a new value or reverts back to
145    an original value.
146
147    @param config_values: See get_moblab_settings().
148
149    @return: The new shadow configuration as ConfigParser object.
150    """
151    original_config = _read_original_config()
152    existing_shadow = _read_raw_config(_CONFIG.shadow_file)
153    for section, config_value_list in config_values.iteritems():
154        for key, value in config_value_list:
155            if original_config.get_config_value(section, key,
156                                                default='',
157                                                allow_blank=True) != value:
158                if not existing_shadow.has_section(section):
159                    existing_shadow.add_section(section)
160                existing_shadow.set(section, key, value)
161            elif existing_shadow.has_option(section, key):
162                existing_shadow.remove_option(section, key)
163    return existing_shadow
164
165
166def _update_partial_config(config_values):
167    """Updates the shadow configuration file with a partial config udpate.
168
169    @param config_values: See get_moblab_settings().
170    """
171    existing_config = _get_shadow_config_from_partial_update(config_values)
172    _write_config_file(_CONFIG.shadow_file, existing_config, True)
173
174
175@rpc_utils.moblab_only
176def update_config_handler(config_values):
177    """Update config values and override shadow config.
178
179    @param config_values: See get_moblab_settings().
180    """
181    original_config = _read_original_config()
182    new_shadow = ConfigParser.RawConfigParser()
183    for section, config_value_list in config_values.iteritems():
184        for key, value in config_value_list:
185            if original_config.get_config_value(section, key,
186                                                default='',
187                                                allow_blank=True) != value:
188                if not new_shadow.has_section(section):
189                    new_shadow.add_section(section)
190                new_shadow.set(section, key, value)
191
192    if not _CONFIG.shadow_file or not os.path.exists(_CONFIG.shadow_file):
193        raise error.RPCException('Shadow config file does not exist.')
194    _write_config_file(_CONFIG.shadow_file, new_shadow, True)
195
196    # TODO (sbasi) crbug.com/403916 - Remove the reboot command and
197    # instead restart the services that rely on the config values.
198    os.system('sudo reboot')
199
200
201@rpc_utils.moblab_only
202def reset_config_settings():
203    """Reset moblab shadow config."""
204    with open(_CONFIG.shadow_file, 'w') as config_file:
205        pass
206    os.system('sudo reboot')
207
208
209@rpc_utils.moblab_only
210def reboot_moblab():
211    """Simply reboot the device."""
212    os.system('sudo reboot')
213
214
215@rpc_utils.moblab_only
216def set_boto_key(boto_key):
217    """Update the boto_key file.
218
219    @param boto_key: File name of boto_key uploaded through handle_file_upload.
220    """
221    if not os.path.exists(boto_key):
222        raise error.RPCException('Boto key: %s does not exist!' % boto_key)
223    shutil.copyfile(boto_key, moblab_host.MOBLAB_BOTO_LOCATION)
224
225
226@rpc_utils.moblab_only
227def set_service_account_credential(service_account_filename):
228    """Update the service account credential file.
229
230    @param service_account_filename: Name of uploaded file through
231            handle_file_upload.
232    """
233    if not os.path.exists(service_account_filename):
234        raise error.RPCException(
235                'Service account file: %s does not exist!' %
236                service_account_filename)
237    shutil.copyfile(
238            service_account_filename,
239            moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION)
240
241
242@rpc_utils.moblab_only
243def set_launch_control_key(launch_control_key):
244    """Update the launch_control_key file.
245
246    @param launch_control_key: File name of launch_control_key uploaded through
247            handle_file_upload.
248    """
249    if not os.path.exists(launch_control_key):
250        raise error.RPCException('Launch Control key: %s does not exist!' %
251                                 launch_control_key)
252    shutil.copyfile(launch_control_key,
253                    moblab_host.MOBLAB_LAUNCH_CONTROL_KEY_LOCATION)
254    # Restart the devserver service.
255    os.system('sudo restart moblab-devserver-init')
256
257
258###########Moblab Config Wizard RPCs #######################
259def _get_public_ip_address(socket_handle):
260    """Gets the public IP address.
261
262    Connects to Google DNS server using a socket and gets the preferred IP
263    address from the connection.
264
265    @param: socket_handle: a unix socket.
266
267    @return: public ip address as string.
268    """
269    try:
270        socket_handle.settimeout(1)
271        socket_handle.connect(('8.8.8.8', 53))
272        socket_name = socket_handle.getsockname()
273        if socket_name is not None:
274            logging.info('Got socket name from UDP socket.')
275            return socket_name[0]
276        logging.warn('Created UDP socket but with no socket_name.')
277    except socket.error:
278        logging.warn('Could not get socket name from UDP socket.')
279    return None
280
281
282def _get_network_info():
283    """Gets the network information.
284
285    TCP socket is used to test the connectivity. If there is no connectivity,
286    try to get the public IP with UDP socket.
287
288    @return: a tuple as (public_ip_address, connected_to_internet).
289    """
290    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
291    ip = _get_public_ip_address(s)
292    if ip is not None:
293        logging.info('Established TCP connection with well known server.')
294        return (ip, True)
295    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
296    return (_get_public_ip_address(s), False)
297
298
299@rpc_utils.moblab_only
300def get_network_info():
301    """Returns the server ip addresses, and if the server connectivity.
302
303    The server ip addresses as an array of strings, and the connectivity as a
304    flag.
305    """
306    network_info = {}
307    info = _get_network_info()
308    if info[0] is not None:
309        network_info['server_ips'] = [info[0]]
310    network_info['is_connected'] = info[1]
311
312    return rpc_utils.prepare_for_serialization(network_info)
313
314
315# Gets the boto configuration.
316def _get_boto_config():
317    """Reads the boto configuration from the boto file.
318
319    @return: Boto configuration as ConfigParser object.
320    """
321    boto_config = ConfigParser.ConfigParser()
322    boto_config.read(MOBLAB_BOTO_LOCATION)
323    return boto_config
324
325
326@rpc_utils.moblab_only
327def get_cloud_storage_info():
328    """RPC handler to get the cloud storage access information.
329    """
330    cloud_storage_info = {}
331    value =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
332    if value is not None:
333        cloud_storage_info[_IMAGE_STORAGE_SERVER] = value
334    value = _CONFIG.get_config_value('CROS', _RESULT_STORAGE_SERVER,
335            default=None)
336    if value is not None:
337        cloud_storage_info[_RESULT_STORAGE_SERVER] = value
338
339    boto_config = _get_boto_config()
340    sections = boto_config.sections()
341
342    if sections:
343        cloud_storage_info[_USE_EXISTING_BOTO_FILE] = True
344    else:
345        cloud_storage_info[_USE_EXISTING_BOTO_FILE] = False
346    if 'Credentials' in sections:
347        options = boto_config.options('Credentials')
348        if _GS_ACCESS_KEY_ID in options:
349            value = boto_config.get('Credentials', _GS_ACCESS_KEY_ID)
350            cloud_storage_info[_GS_ACCESS_KEY_ID] = value
351        if _GS_SECRET_ACCESS_KEY in options:
352            value = boto_config.get('Credentials', _GS_SECRET_ACCESS_KEY)
353            cloud_storage_info[_GS_SECRET_ACCESS_KEY] = value
354
355    return rpc_utils.prepare_for_serialization(cloud_storage_info)
356
357
358def _get_bucket_name_from_url(bucket_url):
359    """Gets the bucket name from a bucket url.
360
361    @param: bucket_url: the bucket url string.
362    """
363    if bucket_url:
364        match = GOOGLE_STORAGE_BUCKET_URL_PATTERN.match(bucket_url)
365        if match:
366            return match.group('bucket')
367    return None
368
369
370def _is_valid_boto_key(key_id, key_secret, directory):
371  try:
372      _run_bucket_performance_test(key_id, key_secret, directory)
373  except BucketPerformanceTestException as e:
374       return(False, str(e))
375  return(True, None)
376
377
378def _validate_cloud_storage_info(cloud_storage_info):
379    """Checks if the cloud storage information is valid.
380
381    @param: cloud_storage_info: The JSON RPC object for cloud storage info.
382
383    @return: A tuple as (valid_boolean, details_string).
384    """
385    valid = True
386    details = None
387    if not cloud_storage_info[_USE_EXISTING_BOTO_FILE]:
388        key_id = cloud_storage_info[_GS_ACCESS_KEY_ID]
389        key_secret = cloud_storage_info[_GS_SECRET_ACCESS_KEY]
390        valid, details = _is_valid_boto_key(
391            key_id, key_secret, cloud_storage_info[_IMAGE_STORAGE_SERVER])
392    return (valid, details)
393
394
395def _create_operation_status_response(is_ok, details):
396    """Helper method to create a operation status reponse.
397
398    @param: is_ok: Boolean for if the operation is ok.
399    @param: details: A detailed string.
400
401    @return: A serialized JSON RPC object.
402    """
403    status_response = {'status_ok': is_ok}
404    if details:
405        status_response['status_details'] = details
406    return rpc_utils.prepare_for_serialization(status_response)
407
408
409@rpc_utils.moblab_only
410def validate_cloud_storage_info(cloud_storage_info):
411    """RPC handler to check if the cloud storage info is valid.
412
413    @param cloud_storage_info: The JSON RPC object for cloud storage info.
414    """
415    valid, details = _validate_cloud_storage_info(cloud_storage_info)
416    return _create_operation_status_response(valid, details)
417
418
419@rpc_utils.moblab_only
420def submit_wizard_config_info(cloud_storage_info, wifi_info):
421    """RPC handler to submit the cloud storage info.
422
423    @param cloud_storage_info: The JSON RPC object for cloud storage info.
424    @param wifi_info: The JSON RPC object for DUT wifi info.
425    """
426    config_update = {}
427    config_update['CROS'] = [
428        (_IMAGE_STORAGE_SERVER, cloud_storage_info[_IMAGE_STORAGE_SERVER]),
429        (_RESULT_STORAGE_SERVER, cloud_storage_info[_RESULT_STORAGE_SERVER])
430    ]
431    config_update['MOBLAB'] = [
432        (_WIFI_AP_NAME, wifi_info.get(_WIFI_AP_NAME) or ''),
433        (_WIFI_AP_PASS, wifi_info.get(_WIFI_AP_PASS) or '')
434    ]
435    _update_partial_config(config_update)
436
437    if not cloud_storage_info[_USE_EXISTING_BOTO_FILE]:
438        boto_config = ConfigParser.RawConfigParser()
439        boto_config.add_section('Credentials')
440        boto_config.set('Credentials', _GS_ACCESS_KEY_ID,
441                        cloud_storage_info[_GS_ACCESS_KEY_ID])
442        boto_config.set('Credentials', _GS_SECRET_ACCESS_KEY,
443                        cloud_storage_info[_GS_SECRET_ACCESS_KEY])
444        _write_config_file(MOBLAB_BOTO_LOCATION, boto_config, True)
445
446    _CONFIG.parse_config_file()
447    _enable_notification_using_credentials_in_bucket()
448    services = ['moblab-devserver-init',
449    'moblab-devserver-cleanup-init', 'moblab-gsoffloader_s-init',
450    'moblab-scheduler-init', 'moblab-gsoffloader-init']
451    cmd = 'export ATEST_RESULTS_DIR=/usr/local/autotest/results;'
452    cmd += 'sudo stop ' + ';sudo stop '.join(services)
453    cmd += ';sudo start ' + ';sudo start '.join(services)
454    cmd += ';sudo apache2 -k graceful'
455    logging.info(cmd)
456    try:
457        utils.run(cmd)
458    except error.CmdError as e:
459        logging.error(e)
460        # if all else fails reboot the device.
461        utils.run('sudo reboot')
462
463    return _create_operation_status_response(True, None)
464
465
466@rpc_utils.moblab_only
467def get_version_info():
468    """ RPC handler to get informaiton about the version of the moblab.
469
470    @return: A serialized JSON RPC object.
471    """
472    lines = open(_ETC_LSB_RELEASE).readlines()
473    version_response = {
474        x.split('=')[0]: x.split('=')[1] for x in lines if '=' in x}
475    version_response['MOBLAB_ID'] = utils.get_moblab_id();
476    version_response['MOBLAB_SERIAL_NUMBER'] = (
477        utils.get_moblab_serial_number())
478    _check_for_system_update()
479    update_status = _get_system_update_status()
480    version_response['MOBLAB_UPDATE_VERSION'] = update_status['NEW_VERSION']
481    version_response['MOBLAB_UPDATE_STATUS'] = update_status['CURRENT_OP']
482    version_response['MOBLAB_UPDATE_PROGRESS'] = update_status['PROGRESS']
483    return rpc_utils.prepare_for_serialization(version_response)
484
485
486@rpc_utils.moblab_only
487def update_moblab():
488    """ RPC call to update and reboot moblab """
489    _install_system_update()
490
491
492def _check_for_system_update():
493    """ Run the ChromeOS update client to check update server for an
494    update. If an update exists, the update client begins downloading it
495    in the background
496    """
497    # sudo is required to run the update client
498    subprocess.call(['sudo', _UPDATE_ENGINE_CLIENT, '--check_for_update'])
499    # wait for update engine to finish checking
500    tries = 0
501    while ('CHECKING_FOR_UPDATE' in _get_system_update_status()['CURRENT_OP']
502            and tries < 10):
503        time.sleep(.1)
504        tries = tries + 1
505
506def _get_system_update_status():
507    """ Run the ChromeOS update client to check status on a
508    pending/downloading update
509
510    @return: A dictionary containing {
511        PROGRESS: str containing percent progress of an update download
512        CURRENT_OP: str current status of the update engine,
513            ex UPDATE_STATUS_UPDATED_NEED_REBOOT
514        NEW_SIZE: str size of the update
515        NEW_VERSION: str version number for the update
516        LAST_CHECKED_TIME: str unix time stamp of the last update check
517    }
518    """
519    # sudo is required to run the update client
520    cmd_out = subprocess.check_output(
521        ['sudo' ,_UPDATE_ENGINE_CLIENT, '--status'])
522    split_lines = [x.split('=') for x in cmd_out.strip().split('\n')]
523    status = dict((key, val) for [key, val] in split_lines)
524    return status
525
526
527def _install_system_update():
528    """ Installs a ChromeOS update, will cause the system to reboot
529    """
530    # sudo is required to run the update client
531    # first run a blocking command to check, fetch, prepare an update
532    # then check if a reboot is needed
533    try:
534        subprocess.check_call(['sudo', _UPDATE_ENGINE_CLIENT, '--update'])
535        # --is_reboot_needed returns 0 if a reboot is required
536        subprocess.check_call(
537            ['sudo', _UPDATE_ENGINE_CLIENT, '--is_reboot_needed'])
538        subprocess.call(['sudo', _UPDATE_ENGINE_CLIENT, '--reboot'])
539
540    except subprocess.CalledProcessError as e:
541        update_error = subprocess.check_output(
542            ['sudo', _UPDATE_ENGINE_CLIENT, '--last_attempt_error'])
543        raise error.RPCException(update_error)
544
545
546@rpc_utils.moblab_only
547def get_connected_dut_info():
548    """ RPC handler to get informaiton about the DUTs connected to the moblab.
549
550    @return: A serialized JSON RPC object.
551    """
552    # Make a list of the connected DUT's
553    leases = _get_dhcp_dut_leases()
554
555
556    connected_duts = _test_all_dut_connections(leases)
557
558    # Get a list of the AFE configured DUT's
559    hosts = list(rpc_utils.get_host_query((), False, True, {}))
560    models.Host.objects.populate_relationships(hosts, models.Label,
561                                               'label_list')
562    configured_duts = {}
563    for host in hosts:
564        labels = [label.name for label in host.label_list]
565        labels.sort()
566        for host_attribute in host.hostattribute_set.all():
567              labels.append("ATTR:(%s=%s)" % (host_attribute.attribute,
568                                              host_attribute.value))
569        configured_duts[host.hostname] = ', '.join(labels)
570
571    return rpc_utils.prepare_for_serialization(
572            {'configured_duts': configured_duts,
573             'connected_duts': connected_duts})
574
575
576def _get_dhcp_dut_leases():
577     """ Extract information about connected duts from the dhcp server.
578
579     @return: A dict of ipaddress to mac address for each device connected.
580     """
581     lease_info = open(_DHCPD_LEASES).read()
582
583     leases = {}
584     for lease in lease_info.split('lease'):
585         if lease.find('binding state active;') != -1:
586             ipaddress = lease.split('\n')[0].strip(' {')
587             last_octet = int(ipaddress.split('.')[-1].strip())
588             if last_octet > 150:
589                 continue
590             mac_address_search = re.search('hardware ethernet (.*);', lease)
591             if mac_address_search:
592                 leases[ipaddress] = mac_address_search.group(1)
593     return leases
594
595def _test_all_dut_connections(leases):
596    """ Test ssh connection of all connected DUTs in parallel
597
598    @param leases: dict containing key value pairs of ip and mac address
599
600    @return: dict containing {
601        ip: {mac_address:[string], ssh_connection_ok:[boolean]}
602    }
603    """
604    # target function for parallel process
605    def _test_dut(ip, result):
606        result.value = _test_dut_ssh_connection(ip)
607
608    processes = []
609    for ip in leases:
610        # use a shared variable to get the ssh test result from child process
611        ssh_test_result = multiprocessing.Value(ctypes.c_bool)
612        # create a subprocess to test each DUT
613        process = multiprocessing.Process(
614            target=_test_dut, args=(ip, ssh_test_result))
615        process.start()
616
617        processes.append({
618            'ip': ip,
619            'ssh_test_result': ssh_test_result,
620            'process': process
621        })
622
623    connected_duts = {}
624    for process in processes:
625        process['process'].join()
626        ip = process['ip']
627        connected_duts[ip] = {
628            'mac_address': leases[ip],
629            'ssh_connection_ok': process['ssh_test_result'].value
630        }
631
632    return connected_duts
633
634
635def _test_dut_ssh_connection(ip):
636    """ Test if a connected dut is accessible via ssh.
637    The primary use case is to verify that the dut has a test image.
638
639    @return: True if the ssh connection is good False else
640    """
641    cmd = ('ssh -o ConnectTimeout=3 -o StrictHostKeyChecking=no '
642            "root@%s 'timeout 2 cat /etc/lsb-release'") % ip
643    try:
644        release = subprocess.check_output(cmd, shell=True)
645        return 'CHROMEOS_RELEASE_APPID' in release
646    except:
647        return False
648
649
650@rpc_utils.moblab_only
651def add_moblab_dut(ipaddress):
652    """ RPC handler to add a connected DUT to autotest.
653
654    @param ipaddress: IP address of the DUT.
655
656    @return: A string giving information about the status.
657    """
658    cmd = '/usr/local/autotest/cli/atest host create %s &' % ipaddress
659    subprocess.call(cmd, shell=True)
660    return (True, 'DUT %s added to Autotest' % ipaddress)
661
662
663@rpc_utils.moblab_only
664def remove_moblab_dut(ipaddress):
665    """ RPC handler to remove DUT entry from autotest.
666
667    @param ipaddress: IP address of the DUT.
668
669    @return: True if the command succeeds without an exception
670    """
671    models.Host.smart_get(ipaddress).delete()
672    return (True, 'DUT %s deleted from Autotest' % ipaddress)
673
674
675@rpc_utils.moblab_only
676def add_moblab_label(ipaddress, label_name):
677    """ RPC handler to add a label in autotest to a DUT entry.
678
679    @param ipaddress: IP address of the DUT.
680    @param label_name: The label name.
681
682    @return: A string giving information about the status.
683    """
684    # Try to create the label in case it does not already exist.
685    label = None
686    try:
687        label = models.Label.add_object(name=label_name)
688    except:
689        label = models.Label.smart_get(label_name)
690        if label.is_replaced_by_static():
691            raise error.UnmodifiableLabelException(
692                    'Failed to add label "%s" because it is a static label. '
693                    'Use go/chromeos-skylab-inventory-tools to add this '
694                    'label.' % label.name)
695
696    host_obj = models.Host.smart_get(ipaddress)
697    if label:
698        label.host_set.add(host_obj)
699        return (True, 'Added label %s to DUT %s' % (label_name, ipaddress))
700    return (False,
701            'Failed to add label %s to DUT %s' % (label_name, ipaddress))
702
703
704@rpc_utils.moblab_only
705def remove_moblab_label(ipaddress, label_name):
706    """ RPC handler to remove a label in autotest from a DUT entry.
707
708    @param ipaddress: IP address of the DUT.
709    @param label_name: The label name.
710
711    @return: A string giving information about the status.
712    """
713    host_obj = models.Host.smart_get(ipaddress)
714    label = models.Label.smart_get(label_name)
715    if label.is_replaced_by_static():
716        raise error.UnmodifiableLabelException(
717                    'Failed to remove label "%s" because it is a static label. '
718                    'Use go/chromeos-skylab-inventory-tools to remove this '
719                    'label.' % label.name)
720
721    label.host_set.remove(host_obj)
722    return (True, 'Removed label %s from DUT %s' % (label_name, ipaddress))
723
724
725@rpc_utils.moblab_only
726def set_host_attrib(ipaddress, attribute, value):
727    """ RPC handler to set an attribute of a host.
728
729    @param ipaddress: IP address of the DUT.
730    @param attribute: string name of attribute
731    @param value: string, or None to delete an attribute
732
733    @return: True if the command succeeds without an exception
734    """
735    host_obj = models.Host.smart_get(ipaddress)
736    host_obj.set_or_delete_attribute(attribute, value)
737    return (True, 'Updated attribute %s to %s on DUT %s' % (
738        attribute, value, ipaddress))
739
740
741@rpc_utils.moblab_only
742def delete_host_attrib(ipaddress, attribute):
743    """ RPC handler to delete an attribute of a host.
744
745    @param ipaddress: IP address of the DUT.
746    @param attribute: string name of attribute
747
748    @return: True if the command succeeds without an exception
749    """
750    host_obj = models.Host.smart_get(ipaddress)
751    host_obj.set_or_delete_attribute(attribute, None)
752    return (True, 'Deleted attribute %s from DUT %s' % (
753        attribute, ipaddress))
754
755
756def _get_connected_dut_labels(requested_label, only_first_label=True):
757    """ Query the DUT's attached to the moblab and return a filtered list
758        of labels.
759
760    @param requested_label:  the label name you are requesting.
761    @param only_first_label:  if the device has the same label name multiple
762                              times only return the first label value in the
763                              list.
764
765    @return: A de-duped list of requested dut labels attached to the moblab.
766    """
767    hosts = list(rpc_utils.get_host_query((), False, True, {}))
768    if not hosts:
769        return []
770    models.Host.objects.populate_relationships(hosts, models.Label,
771                                               'label_list')
772    labels = set()
773    for host in hosts:
774        for label in host.label_list:
775            if requested_label in label.name:
776                labels.add(label.name.replace(requested_label, ''))
777                if only_first_label:
778                    break
779    return list(labels)
780
781def _get_connected_dut_board_models():
782    """ Get the boards and their models of attached DUTs
783
784    @return: A de-duped list of dut board/model attached to the moblab
785    format: [
786        {
787            "board": "carl",
788            "model": "bruce"
789        },
790        {
791            "board": "veyron_minnie",
792            "model": "veyron_minnie"
793        }
794    ]
795    """
796    hosts = list(rpc_utils.get_host_query((), False, True, {}))
797    if not hosts:
798        return []
799    models.Host.objects.populate_relationships(hosts, models.Label,
800                                               'label_list')
801    model_board_map = dict()
802    for host in hosts:
803        model = ''
804        board = ''
805        for label in host.label_list:
806            if 'model:' in label.name:
807                model = label.name.replace('model:', '')
808            elif 'board:' in label.name:
809                board = label.name.replace('board:', '')
810        model_board_map[model] = board
811
812    board_models_list = []
813    for model in sorted(model_board_map.keys()):
814        board_models_list.append({
815            'model': model,
816            'board': model_board_map[model]
817        })
818    return board_models_list
819
820
821@rpc_utils.moblab_only
822def get_connected_boards():
823    """ RPC handler to get a list of the boards connected to the moblab.
824
825    @return: A de-duped list of board types attached to the moblab.
826    """
827    return _get_connected_dut_board_models()
828
829
830@rpc_utils.moblab_only
831def get_connected_pools():
832    """ RPC handler to get a list of the pools labels on the DUT's connected.
833
834    @return: A de-duped list of pool labels.
835    """
836    pools = _get_connected_dut_labels("pool:", False)
837    pools.sort()
838    return pools
839
840
841@rpc_utils.moblab_only
842def get_builds_for_board(board_name):
843    """ RPC handler to find the most recent builds for a board.
844
845
846    @param board_name: The name of a connected board.
847    @return: A list of string with the most recent builds for the latest
848             three milestones.
849    """
850    return _get_builds_for_in_directory(board_name + '-release',
851                                        milestone_limit=4)
852
853
854@rpc_utils.moblab_only
855def get_firmware_for_board(board_name):
856    """ RPC handler to find the most recent firmware for a board.
857
858
859    @param board_name: The name of a connected board.
860    @return: A list of strings with the most recent firmware builds for the
861             latest three milestones.
862    """
863    return _get_builds_for_in_directory(board_name + '-firmware')
864
865
866def _get_sortable_build_number(sort_key):
867    """ Converts a build number line cyan-release/R59-9460.27.0 into an integer.
868
869        To be able to sort a list of builds you need to convert the build number
870        into an integer so it can be compared correctly to other build.
871
872        cyan-release/R59-9460.27.0 =>  5909460027000
873
874        If the sort key is not recognised as a build number 1 will be returned.
875
876    @param sort_key: A string that represents a build number like
877                     cyan-release/R59-9460.27.0
878    @return: An integer that represents that build number or 1 if not recognised
879             as a build.
880    """
881    build_number = re.search('.*/R([0-9]*)-([0-9]*)\.([0-9]*)\.([0-9]*)',
882                             sort_key)
883    if not build_number or not len(build_number.groups()) == 4:
884      return 1
885    return int("%d%05d%03d%03d" % (int(build_number.group(1)),
886                                   int(build_number.group(2)),
887                                   int(build_number.group(3)),
888                                   int(build_number.group(4))))
889
890def _get_builds_for_in_directory(directory_name, milestone_limit=3,
891                                 build_limit=20):
892    """ Fetch the most recent builds for the last three milestones from gcs.
893
894
895    @param directory_name: The sub-directory under the configured GCS image
896                           storage bucket to search.
897
898
899    @return: A string list no longer than <milestone_limit> x <build_limit>
900             items, containing the most recent <build_limit> builds from the
901             last milestone_limit milestones.
902    """
903    output = StringIO.StringIO()
904    gs_image_location =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
905    try:
906        utils.run(GsUtil.get_gsutil_cmd(),
907                  args=('ls', gs_image_location + directory_name),
908                  stdout_tee=output)
909    except error.CmdError as e:
910        error_text = ('Failed to list builds from %s.\n'
911                'Did you configure your boto key? Try running the config '
912                'wizard again.\n\n%s') % ((gs_image_location + directory_name),
913                    e.result_obj.stderr)
914        raise error.RPCException(error_text)
915    lines = output.getvalue().split('\n')
916    output.close()
917    builds = [line.replace(gs_image_location,'').strip('/ ')
918              for line in lines if line != '']
919    build_matcher = re.compile(r'^.*\/R([0-9]*)-.*')
920    build_map = {}
921    for build in builds:
922        match = build_matcher.match(build)
923        if match:
924            milestone = match.group(1)
925            if milestone not in build_map:
926                build_map[milestone] = []
927            build_map[milestone].append(build)
928    milestones = build_map.keys()
929    milestones.sort()
930    milestones.reverse()
931    build_list = []
932    for milestone in milestones[:milestone_limit]:
933         builds = build_map[milestone]
934         builds.sort(key=_get_sortable_build_number)
935         builds.reverse()
936         build_list.extend(builds[:build_limit])
937    return build_list
938
939
940def _run_bucket_performance_test(key_id, key_secret, bucket_name,
941                                 test_size='1M', iterations='1',
942                                 result_file='/tmp/gsutil_perf.json'):
943    """Run a gsutil perfdiag on a supplied bucket and output the results"
944
945       @param key_id: boto key of the bucket to be accessed
946       @param key_secret: boto secret of the bucket to be accessed
947       @param bucket_name: bucket to be tested.
948       @param test_size: size of file to use in test, see gsutil perfdiag help.
949       @param iterations: number of times each test is run.
950       @param result_file: name of file to write results out to.
951
952       @return None
953       @raises BucketPerformanceTestException if the command fails.
954    """
955    try:
956      utils.run(GsUtil.get_gsutil_cmd(), args=(
957          '-o', 'Credentials:gs_access_key_id=%s' % key_id,
958          '-o', 'Credentials:gs_secret_access_key=%s' % key_secret,
959          'perfdiag', '-s', test_size, '-o', result_file,
960          '-n', iterations,
961          bucket_name))
962    except error.CmdError as e:
963       logging.error(e)
964       # Extract useful error from the stacktrace
965       errormsg = str(e)
966       start_error_pos = errormsg.find("<Error>")
967       end_error_pos = errormsg.find("</Error>", start_error_pos)
968       extracted_error_msg = errormsg[start_error_pos:end_error_pos]
969       raise BucketPerformanceTestException(
970           extracted_error_msg if extracted_error_msg else errormsg)
971    # TODO(haddowk) send the results to the cloud console when that feature is
972    # enabled.
973
974
975# TODO(haddowk) Change suite_args name to "test_filter_list" or similar. May
976# also need to make changes at MoblabRpcHelper.java
977@rpc_utils.moblab_only
978def run_suite(board, build, suite, model=None, ro_firmware=None,
979              rw_firmware=None, pool=None, suite_args=None, test_args=None,
980              bug_id=None, part_id=None):
981    """ RPC handler to run a test suite.
982
983    @param board: a board name connected to the moblab.
984    @param build: a build name of a build in the GCS.
985    @param suite: the name of a suite to run
986    @param model: a board model name connected to the moblab.
987    @param ro_firmware: Optional ro firmware build number to use.
988    @param rw_firmware: Optional rw firmware build number to use.
989    @param pool: Optional pool name to run the suite in.
990    @param suite_args: Arguments to be used in the suite control file.
991    @param test_args: '\n' delimited key=val pairs passed to test control file.
992    @param bug_id: Optional bug ID used for AVL qualification process.
993    @param part_id: Optional part ID used for AVL qualification
994    process.
995
996    @return: None
997    """
998    builds = {'cros-version': build}
999    # TODO(mattmallett b/92031054) Standardize bug id, part id passing for memory/storage qual
1000    processed_suite_args = dict()
1001    processed_test_args = dict()
1002    if rw_firmware:
1003        builds['fwrw-version'] = rw_firmware
1004    if ro_firmware:
1005        builds['fwro-version'] = ro_firmware
1006    if suite_args:
1007        processed_suite_args['tests'] = \
1008            [s.strip() for s in suite_args.split(',')]
1009    if bug_id:
1010        processed_suite_args['bug_id'] = bug_id
1011    if part_id:
1012        processed_suite_args['part_id'] = part_id
1013    processed_test_args['bug_id'] = bug_id or ''
1014    processed_test_args['part_id'] = part_id or ''
1015
1016
1017    # set processed_suite_args to None instead of empty dict when there is no
1018    # argument in processed_suite_args
1019    if len(processed_suite_args) == 0:
1020        processed_suite_args = None
1021
1022    if test_args:
1023        try:
1024          processed_test_args['args'] = [test_args]
1025          for line in test_args.split('\n'):
1026              key, value = line.strip().split('=')
1027              processed_test_args[key] = value
1028        except:
1029            raise error.RPCException('Could not parse test args.')
1030
1031
1032    ap_name =_CONFIG.get_config_value('MOBLAB', _WIFI_AP_NAME, default=None)
1033    processed_test_args['ssid'] = ap_name
1034    ap_pass =_CONFIG.get_config_value('MOBLAB', _WIFI_AP_PASS, default='')
1035    processed_test_args['wifipass'] = ap_pass
1036
1037    suite_timeout_mins = _SUITE_TIMEOUT_MAP.get(
1038            suite, _DEFAULT_SUITE_TIMEOUT_MINS)
1039
1040    afe = frontend.AFE(user='moblab')
1041    afe.run('create_suite_job', board=board, builds=builds, name=suite,
1042            pool=pool, run_prod_code=False, test_source_build=build,
1043            wait_for_results=True, suite_args=processed_suite_args,
1044            test_args=processed_test_args, job_retry=True,
1045            max_retries=sys.maxint, model=model,
1046            timeout_mins=suite_timeout_mins,
1047            max_runtime_mins=suite_timeout_mins)
1048
1049
1050def _enable_notification_using_credentials_in_bucket():
1051    """ Check and enable cloud notification if a credentials file exits.
1052    @return: None
1053    """
1054    gs_image_location =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
1055    try:
1056        utils.run(GsUtil.get_gsutil_cmd(), args=(
1057            'cp', gs_image_location + 'pubsub-key-do-not-delete.json', '/tmp'))
1058        # This runs the copy as moblab user
1059        shutil.copyfile('/tmp/pubsub-key-do-not-delete.json',
1060                        moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION)
1061
1062    except error.CmdError as e:
1063        logging.error(e)
1064    else:
1065        logging.info('Enabling cloud notifications')
1066        config_update = {}
1067        config_update['CROS'] = [(_CLOUD_NOTIFICATION_ENABLED, True)]
1068        _update_partial_config(config_update)
1069
1070
1071@rpc_utils.moblab_only
1072def get_dut_wifi_info():
1073    """RPC handler to get the dut wifi AP information.
1074    """
1075    dut_wifi_info = {}
1076    value =_CONFIG.get_config_value('MOBLAB', _WIFI_AP_NAME,
1077        default=None)
1078    if value is not None:
1079        dut_wifi_info[_WIFI_AP_NAME] = value
1080    value = _CONFIG.get_config_value('MOBLAB', _WIFI_AP_PASS,
1081        default=None)
1082    if value is not None:
1083        dut_wifi_info[_WIFI_AP_PASS] = value
1084    return rpc_utils.prepare_for_serialization(dut_wifi_info)
1085