1#!/usr/bin/env python2 2 3# Copyright 2015 The Chromium OS Authors. All rights reserved. 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7"""Tool to sync lab servers to the "Allowed Networks" of a CloudSQL instance. 8 9For a lab server to access CloudSQL instance, the server's IP must be added to 10the "Allowed Networks" list of the CloudSQL instance. This tool is to be used to 11read the list of lab servers from server database and update the list of 12"Allowed Networks" of a given CloudSQL instance. 13 14The tool also reads CLOUD/tko_access_servers from global config to add these 15servers to the "Allowed Networks" list of the CloudSQL instance. This allows 16servers that do not run Autotest code can access the CloudSQL instance. 17 18Note that running this tool will overwrite existing IPs in the "Allowed 19Networks" list. Therefore, manually editing that list from CloudSQL console 20should be prohibited. Instead, the servers should be added to 21CLOUD/tko_access_servers in shadow_config.ini. 22 23""" 24 25import argparse 26import socket 27import sys 28 29import common 30from autotest_lib.client.bin import utils 31from autotest_lib.client.common_lib import error 32from autotest_lib.client.common_lib import global_config 33from autotest_lib.client.common_lib.cros import retry 34from autotest_lib.server import frontend 35 36 37ROLES_REQUIRE_TKO_ACCESS = { 38 'afe', 39 'database', 40 'drone', 41 'scheduler', 42 'sentinel', 43 'shard', 44 'skylab_drone', 45} 46 47 48def gcloud_login(project): 49 """Login to Google Cloud service for gcloud command to run. 50 51 @param project: Name of the Google Cloud project. 52 """ 53 # Login with user account. If the user hasn't log in yet, the script will 54 # print a url and ask for a verification code. User should load the url in 55 # browser, and copy the verification code from the web page. When private IP 56 # can be supported to be added using non-corp account, the login can be done 57 # through service account and key file, e.g., 58 # gcloud auth activate-service-account --key-file ~/key.json 59 utils.run('gcloud auth login', stdout_tee=sys.stdout, 60 stderr_tee=sys.stderr, stdin=sys.stdin) 61 62 63@retry.retry(error.CmdError, timeout_min=3) 64def _fetch_external_ip(server_name): 65 return utils.run('ssh %s curl -s ifconfig.me' % server_name).stdout.rstrip() 66 67 68def update_allowed_networks(project, instance, afe=None, extra_servers=None, 69 dryrun=False): 70 """Update the "Allowed Networks" list of the given CloudSQL instance. 71 72 @param project: Name of the Google Cloud project. 73 @param instance: Name of the CloudSQL instance. 74 @param afe: Server of the frontend RPC, default to None to use the server 75 specified in global config. 76 @param extra_servers: Extra servers to be included in the "Allowed Networks" 77 list. Default is None. 78 @param dryrun: Boolean indicating whether this is a dryrun. 79 """ 80 # Get the IP address of all servers need access to CloudSQL instance. 81 rpc = frontend.AFE(server=afe) 82 servers = [s['hostname'] for s in rpc.run('get_servers') 83 if s['status'] != 'repair_required' and 84 ROLES_REQUIRE_TKO_ACCESS.intersection(s['roles'])] 85 if extra_servers: 86 servers.extend(extra_servers.split(',')) 87 # Extra servers can be listed in CLOUD/tko_access_servers shadow config. 88 tko_servers = global_config.global_config.get_config_value( 89 'CLOUD', 'tko_access_servers', default='') 90 if tko_servers: 91 servers.extend(tko_servers.split(',')) 92 print 'Adding servers %s to access list for projects %s' % (servers, 93 instance) 94 print 'Fetching their IP addresses...' 95 ips = [] 96 for name in servers: 97 try: 98 # collect internal ips 99 ips.append(socket.gethostbyname(name)) 100 # collect external ips 101 ips.append(_fetch_external_ip(name)) 102 except socket.gaierror: 103 print 'Failed to resolve internal IP address for name %s' % name 104 raise 105 except error.TimeoutException: 106 print 'Failed to resolve external IP address for %s' % name 107 raise 108 109 print '...Done: %s' % ips 110 111 cidr_ips = [str(ip) + '/32' for ip in ips] 112 113 if dryrun: 114 print 'This is a dryrun: skip updating glcoud sql whitelists.' 115 return 116 117 login = False 118 while True: 119 try: 120 utils.run('gcloud config set project %s -q' % project) 121 cmd = ('gcloud sql instances patch %s --authorized-networks %s ' 122 '-q' % (instance, ','.join(cidr_ips))) 123 print 'Running command to update whitelists: "%s"' % cmd 124 utils.run(cmd, stdout_tee=sys.stdout, stderr_tee=sys.stderr) 125 return 126 except error.CmdError: 127 if login: 128 raise 129 130 # Try to login and retry if the command failed. 131 gcloud_login(project) 132 login = True 133 134 135def main(): 136 """main script.""" 137 parser = argparse.ArgumentParser() 138 parser.add_argument('--project', type=str, dest='project', 139 help='Name of the Google Cloud project.') 140 parser.add_argument('--instance', type=str, dest='instance', 141 help='Name of the CloudSQL instance.') 142 parser.add_argument('--afe', type=str, dest='afe', 143 help='Name of the RPC server to get server list.', 144 default=None) 145 parser.add_argument('--extra_servers', type=str, dest='extra_servers', 146 help=('Extra servers to be included in the "Allowed ' 147 'Networks" list separated by comma.'), 148 default=None) 149 parser.add_argument('--dryrun', dest='dryrun', action='store_true', 150 default=False, 151 help='Fetch IPs without updating whitelists in gcloud.') 152 options = parser.parse_args() 153 154 update_allowed_networks(options.project, options.instance, options.afe, 155 options.extra_servers, options.dryrun) 156 157 158if __name__ == '__main__': 159 main() 160