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