Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 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 | |
Prathmesh Prabhu | aaa5586 | 2018-01-31 10:06:22 -0800 | [diff] [blame] | 36 | ROLES_REQUIRE_TKO_ACCESS = { |
| 37 | 'afe', |
| 38 | 'database', |
| 39 | 'drone', |
| 40 | 'scheduler', |
| 41 | 'sentinel', |
| 42 | 'shard', |
| 43 | } |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 44 | |
| 45 | def gcloud_login(project): |
| 46 | """Login to Google Cloud service for gcloud command to run. |
| 47 | |
| 48 | @param project: Name of the Google Cloud project. |
| 49 | """ |
| 50 | # Login with user account. If the user hasn't log in yet, the script will |
| 51 | # print a url and ask for a verification code. User should load the url in |
| 52 | # browser, and copy the verification code from the web page. When private IP |
| 53 | # can be supported to be added using non-corp account, the login can be done |
| 54 | # through service account and key file, e.g., |
| 55 | # gcloud auth activate-service-account --key-file ~/key.json |
| 56 | utils.run('gcloud auth login', stdout_tee=sys.stdout, |
| 57 | stderr_tee=sys.stderr, stdin=sys.stdin) |
| 58 | |
| 59 | |
| 60 | def update_allowed_networks(project, instance, afe=None, extra_servers=None): |
| 61 | """Update the "Allowed Networks" list of the given CloudSQL instance. |
| 62 | |
| 63 | @param project: Name of the Google Cloud project. |
| 64 | @param instance: Name of the CloudSQL instance. |
| 65 | @param afe: Server of the frontend RPC, default to None to use the server |
| 66 | specified in global config. |
| 67 | @param extra_servers: Extra servers to be included in the "Allowed Networks" |
| 68 | list. Default is None. |
| 69 | """ |
| 70 | # Get the IP address of all servers need access to CloudSQL instance. |
| 71 | rpc = frontend.AFE(server=afe) |
| 72 | servers = [s['hostname'] for s in rpc.run('get_servers') |
| 73 | if s['status'] != 'repair_required' and |
| 74 | ROLES_REQUIRE_TKO_ACCESS.intersection(s['roles'])] |
| 75 | if extra_servers: |
| 76 | servers.extend(extra_servers.split(',')) |
| 77 | # Extra servers can be listed in CLOUD/tko_access_servers shadow config. |
| 78 | tko_servers = global_config.global_config.get_config_value( |
| 79 | 'CLOUD', 'tko_access_servers', default='') |
| 80 | if tko_servers: |
| 81 | servers.extend(tko_servers.split(',')) |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 82 | print 'Adding servers %s to access list for projects %s' % (servers, |
| 83 | instance) |
| 84 | print 'Fetching their IP addresses...' |
Aviv Keshet | 4eb2a89 | 2018-02-06 16:51:18 -0800 | [diff] [blame^] | 85 | ips = [] |
| 86 | for name in servers: |
| 87 | try: |
| 88 | ips.append(socket.gethostbyname(name)) |
| 89 | except socket.gaierror: |
| 90 | print 'Failed to resolve IP address for name %s' % name |
| 91 | raise |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 92 | print '...Done: %s' % ips |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 93 | |
| 94 | login = False |
| 95 | while True: |
| 96 | try: |
Aviv Keshet | 0495e15 | 2017-05-01 19:31:44 -0700 | [diff] [blame] | 97 | utils.run('gcloud config set project %s -q' % project) |
| 98 | cmd = ('gcloud sql instances patch %s --authorized-networks %s ' |
| 99 | '-q' % (instance, ','.join(ips))) |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 100 | print 'Running command to update whitelists: "%s"' % cmd |
| 101 | utils.run(cmd, stdout_tee=sys.stdout, stderr_tee=sys.stderr) |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 102 | return |
| 103 | except error.CmdError: |
| 104 | if login: |
| 105 | raise |
| 106 | |
| 107 | # Try to login and retry if the command failed. |
| 108 | gcloud_login(project) |
| 109 | login = True |
| 110 | |
| 111 | |
| 112 | def main(): |
| 113 | """main script.""" |
| 114 | parser = argparse.ArgumentParser() |
| 115 | parser.add_argument('--project', type=str, dest='project', |
| 116 | help='Name of the Google Cloud project.') |
| 117 | parser.add_argument('--instance', type=str, dest='instance', |
| 118 | help='Name of the CloudSQL instance.') |
| 119 | parser.add_argument('--afe', type=str, dest='afe', |
| 120 | help='Name of the RPC server to get server list.', |
| 121 | default=None) |
| 122 | parser.add_argument('--extra_servers', type=str, dest='extra_servers', |
| 123 | help=('Extra servers to be included in the "Allowed ' |
| 124 | 'Networks" list separated by comma.'), |
| 125 | default=None) |
| 126 | options = parser.parse_args() |
| 127 | |
| 128 | update_allowed_networks(options.project, options.instance, options.afe, |
| 129 | options.extra_servers) |
| 130 | |
| 131 | |
| 132 | if __name__ == '__main__': |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 133 | main() |