Mike Frysinger | d03e6b5 | 2019-08-03 12:49:01 -0400 | [diff] [blame] | 1 | #!/usr/bin/env python2 |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 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 |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 33 | from autotest_lib.client.common_lib.cros import retry |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 34 | from autotest_lib.server import frontend |
| 35 | |
| 36 | |
Prathmesh Prabhu | aaa5586 | 2018-01-31 10:06:22 -0800 | [diff] [blame] | 37 | ROLES_REQUIRE_TKO_ACCESS = { |
| 38 | 'afe', |
| 39 | 'database', |
| 40 | 'drone', |
| 41 | 'scheduler', |
| 42 | 'sentinel', |
| 43 | 'shard', |
Prathmesh Prabhu | 5ce9489 | 2018-07-26 20:38:22 -0700 | [diff] [blame] | 44 | 'skylab_drone', |
Prathmesh Prabhu | aaa5586 | 2018-01-31 10:06:22 -0800 | [diff] [blame] | 45 | } |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 46 | |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 47 | |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 48 | def 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 | |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 63 | @retry.retry(error.CmdError, timeout_min=3) |
| 64 | def _fetch_external_ip(server_name): |
| 65 | return utils.run('ssh %s curl -s ifconfig.me' % server_name).stdout.rstrip() |
| 66 | |
| 67 | |
| 68 | def update_allowed_networks(project, instance, afe=None, extra_servers=None, |
| 69 | dryrun=False): |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 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. |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 78 | @param dryrun: Boolean indicating whether this is a dryrun. |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 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(',')) |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 92 | print 'Adding servers %s to access list for projects %s' % (servers, |
| 93 | instance) |
| 94 | print 'Fetching their IP addresses...' |
Aviv Keshet | 4eb2a89 | 2018-02-06 16:51:18 -0800 | [diff] [blame] | 95 | ips = [] |
| 96 | for name in servers: |
| 97 | try: |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 98 | # collect internal ips |
Aviv Keshet | 4eb2a89 | 2018-02-06 16:51:18 -0800 | [diff] [blame] | 99 | ips.append(socket.gethostbyname(name)) |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 100 | # collect external ips |
| 101 | ips.append(_fetch_external_ip(name)) |
Aviv Keshet | 4eb2a89 | 2018-02-06 16:51:18 -0800 | [diff] [blame] | 102 | except socket.gaierror: |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 103 | print 'Failed to resolve internal IP address for name %s' % name |
Aviv Keshet | 4eb2a89 | 2018-02-06 16:51:18 -0800 | [diff] [blame] | 104 | raise |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 105 | except error.TimeoutException: |
| 106 | print 'Failed to resolve external IP address for %s' % name |
| 107 | raise |
| 108 | |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 109 | print '...Done: %s' % ips |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 110 | |
Aviv Keshet | fd793bd | 2018-02-12 15:29:35 -0800 | [diff] [blame] | 111 | cidr_ips = [str(ip) + '/32' for ip in ips] |
| 112 | |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 113 | if dryrun: |
| 114 | print 'This is a dryrun: skip updating glcoud sql whitelists.' |
| 115 | return |
| 116 | |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 117 | login = False |
| 118 | while True: |
| 119 | try: |
Aviv Keshet | 0495e15 | 2017-05-01 19:31:44 -0700 | [diff] [blame] | 120 | utils.run('gcloud config set project %s -q' % project) |
| 121 | cmd = ('gcloud sql instances patch %s --authorized-networks %s ' |
Aviv Keshet | fd793bd | 2018-02-12 15:29:35 -0800 | [diff] [blame] | 122 | '-q' % (instance, ','.join(cidr_ips))) |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 123 | print 'Running command to update whitelists: "%s"' % cmd |
| 124 | utils.run(cmd, stdout_tee=sys.stdout, stderr_tee=sys.stderr) |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 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 | |
| 135 | def 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) |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 149 | parser.add_argument('--dryrun', dest='dryrun', action='store_true', |
| 150 | default=False, |
| 151 | help='Fetch IPs without updating whitelists in gcloud.') |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 152 | options = parser.parse_args() |
| 153 | |
| 154 | update_allowed_networks(options.project, options.instance, options.afe, |
Ningning Xia | aeb8b63 | 2018-05-03 12:14:58 -0700 | [diff] [blame] | 155 | options.extra_servers, options.dryrun) |
Dan Shi | a2fa8d1 | 2015-07-22 11:49:38 -0700 | [diff] [blame] | 156 | |
| 157 | |
| 158 | if __name__ == '__main__': |
Aviv Keshet | 2eea6b3 | 2017-05-01 18:21:04 -0700 | [diff] [blame] | 159 | main() |