1 #!/usr/bin/env python 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 9 For a lab server to access CloudSQL instance, the server's IP must be added to 10 the "Allowed Networks" list of the CloudSQL instance. This tool is to be used to 11 read the list of lab servers from server database and update the list of 12 "Allowed Networks" of a given CloudSQL instance. 13 14 The tool also reads CLOUD/tko_access_servers from global config to add these 15 servers to the "Allowed Networks" list of the CloudSQL instance. This allows 16 servers that do not run Autotest code can access the CloudSQL instance. 17 18 Note that running this tool will overwrite existing IPs in the "Allowed 19 Networks" list. Therefore, manually editing that list from CloudSQL console 20 should be prohibited. Instead, the servers should be added to 21 CLOUD/tko_access_servers in shadow_config.ini. 22 23 """ 24 25 import argparse 26 import socket 27 import sys 28 29 import common 30 from autotest_lib.client.bin import utils 31 from autotest_lib.client.common_lib import error 32 from autotest_lib.client.common_lib import global_config 33 from autotest_lib.server import frontend 34 35 36 ROLES_REQUIRE_TKO_ACCESS = {'scheduler', 'drone', 'shard', 'database', 'afe'} 37 38 def gcloud_login(project): 39 """Login to Google Cloud service for gcloud command to run. 40 41 @param project: Name of the Google Cloud project. 42 """ 43 # Login with user account. If the user hasn't log in yet, the script will 44 # print a url and ask for a verification code. User should load the url in 45 # browser, and copy the verification code from the web page. When private IP 46 # can be supported to be added using non-corp account, the login can be done 47 # through service account and key file, e.g., 48 # gcloud auth activate-service-account --key-file ~/key.json 49 utils.run('gcloud auth login', stdout_tee=sys.stdout, 50 stderr_tee=sys.stderr, stdin=sys.stdin) 51 52 53 def update_allowed_networks(project, instance, afe=None, extra_servers=None): 54 """Update the "Allowed Networks" list of the given CloudSQL instance. 55 56 @param project: Name of the Google Cloud project. 57 @param instance: Name of the CloudSQL instance. 58 @param afe: Server of the frontend RPC, default to None to use the server 59 specified in global config. 60 @param extra_servers: Extra servers to be included in the "Allowed Networks" 61 list. Default is None. 62 """ 63 # Get the IP address of all servers need access to CloudSQL instance. 64 rpc = frontend.AFE(server=afe) 65 servers = [s['hostname'] for s in rpc.run('get_servers') 66 if s['status'] != 'repair_required' and 67 ROLES_REQUIRE_TKO_ACCESS.intersection(s['roles'])] 68 if extra_servers: 69 servers.extend(extra_servers.split(',')) 70 # Extra servers can be listed in CLOUD/tko_access_servers shadow config. 71 tko_servers = global_config.global_config.get_config_value( 72 'CLOUD', 'tko_access_servers', default='') 73 if tko_servers: 74 servers.extend(tko_servers.split(',')) 75 ips = [socket.gethostbyname(name) for name in servers] 76 77 login = False 78 while True: 79 try: 80 utils.run('gcloud config set project %s' % project) 81 utils.run('gcloud sql instances patch %s --authorized-networks %s' 82 % (instance, ','.join(ips)), stdout_tee=sys.stdout, 83 stderr_tee=sys.stderr) 84 return 85 except error.CmdError: 86 if login: 87 raise 88 89 # Try to login and retry if the command failed. 90 gcloud_login(project) 91 login = True 92 93 94 def main(): 95 """main script.""" 96 parser = argparse.ArgumentParser() 97 parser.add_argument('--project', type=str, dest='project', 98 help='Name of the Google Cloud project.') 99 parser.add_argument('--instance', type=str, dest='instance', 100 help='Name of the CloudSQL instance.') 101 parser.add_argument('--afe', type=str, dest='afe', 102 help='Name of the RPC server to get server list.', 103 default=None) 104 parser.add_argument('--extra_servers', type=str, dest='extra_servers', 105 help=('Extra servers to be included in the "Allowed ' 106 'Networks" list separated by comma.'), 107 default=None) 108 options = parser.parse_args() 109 110 update_allowed_networks(options.project, options.instance, options.afe, 111 options.extra_servers) 112 113 114 if __name__ == '__main__': 115 main()
