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