• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
25from __future__ import absolute_import
26from __future__ import division
27from __future__ import print_function
28
29import argparse
30import socket
31import sys
32
33import common
34from autotest_lib.client.bin import utils
35from autotest_lib.client.common_lib import error
36from autotest_lib.client.common_lib import global_config
37from autotest_lib.client.common_lib.cros import retry
38from autotest_lib.server import frontend
39
40
41ROLES_REQUIRE_TKO_ACCESS = {
42        'afe',
43        'database',
44        'drone',
45        'scheduler',
46        'sentinel',
47        'shard',
48        'skylab_drone',
49}
50
51
52def gcloud_login(project):
53    """Login to Google Cloud service for gcloud command to run.
54
55    @param project: Name of the Google Cloud project.
56    """
57    # Login with user account. If the user hasn't log in yet, the script will
58    # print a url and ask for a verification code. User should load the url in
59    # browser, and copy the verification code from the web page. When private IP
60    # can be supported to be added using non-corp account, the login can be done
61    # through service account and key file, e.g.,
62    # gcloud auth activate-service-account --key-file ~/key.json
63    utils.run('gcloud auth login', stdout_tee=sys.stdout,
64              stderr_tee=sys.stderr, stdin=sys.stdin)
65
66
67@retry.retry(error.CmdError, timeout_min=3)
68def _fetch_external_ip(server_name):
69    return utils.run('ssh %s curl -s ifconfig.me' % server_name).stdout.rstrip()
70
71
72def update_allowed_networks(project, instance, afe=None, extra_servers=None,
73                            dryrun=False):
74    """Update the "Allowed Networks" list of the given CloudSQL instance.
75
76    @param project: Name of the Google Cloud project.
77    @param instance: Name of the CloudSQL instance.
78    @param afe: Server of the frontend RPC, default to None to use the server
79                specified in global config.
80    @param extra_servers: Extra servers to be included in the "Allowed Networks"
81                          list. Default is None.
82    @param dryrun: Boolean indicating whether this is a dryrun.
83    """
84    # Get the IP address of all servers need access to CloudSQL instance.
85    rpc = frontend.AFE(server=afe)
86    servers = [s['hostname'] for s in rpc.run('get_servers')
87               if s['status'] != 'repair_required' and
88               ROLES_REQUIRE_TKO_ACCESS.intersection(s['roles'])]
89    if extra_servers:
90        servers.extend(extra_servers.split(','))
91    # Extra servers can be listed in CLOUD/tko_access_servers shadow config.
92    tko_servers = global_config.global_config.get_config_value(
93            'CLOUD', 'tko_access_servers', default='')
94    if tko_servers:
95        servers.extend(tko_servers.split(','))
96    print('Adding servers %s to access list for projects %s' % (servers,
97                                                                instance))
98    print('Fetching their IP addresses...')
99    ips = []
100    for name in servers:
101        try:
102            # collect internal ips
103            ips.append(socket.gethostbyname(name))
104            # collect external ips
105            ips.append(_fetch_external_ip(name))
106        except socket.gaierror:
107            print('Failed to resolve internal IP address for name %s' % name)
108            raise
109        except error.TimeoutException:
110            print('Failed to resolve external IP address for %s' % name)
111            raise
112
113    print('...Done: %s' % ips)
114
115    cidr_ips = [str(ip) + '/32' for ip in ips]
116
117    if dryrun:
118        print('This is a dryrun: skip updating glcoud sql allowlists.')
119        return
120
121    login = False
122    while True:
123        try:
124            utils.run('gcloud config set project %s -q' % project)
125            cmd = ('gcloud sql instances patch %s --authorized-networks %s '
126                   '-q' % (instance, ','.join(cidr_ips)))
127            print('Running command to update allowlists: "%s"' % cmd)
128            utils.run(cmd, stdout_tee=sys.stdout, stderr_tee=sys.stderr)
129            return
130        except error.CmdError:
131            if login:
132                raise
133
134            # Try to login and retry if the command failed.
135            gcloud_login(project)
136            login = True
137
138
139def main():
140    """main script."""
141    parser = argparse.ArgumentParser()
142    parser.add_argument('--project', type=str, dest='project',
143                        help='Name of the Google Cloud project.')
144    parser.add_argument('--instance', type=str, dest='instance',
145                        help='Name of the CloudSQL instance.')
146    parser.add_argument('--afe', type=str, dest='afe',
147                        help='Name of the RPC server to get server list.',
148                        default=None)
149    parser.add_argument('--extra_servers', type=str, dest='extra_servers',
150                        help=('Extra servers to be included in the "Allowed '
151                              'Networks" list separated by comma.'),
152                        default=None)
153    parser.add_argument('--dryrun', dest='dryrun', action='store_true',
154                        default=False,
155                        help='Fetch IPs without updating allowlists in gcloud.')
156    options = parser.parse_args()
157
158    update_allowed_networks(options.project, options.instance, options.afe,
159                            options.extra_servers, options.dryrun)
160
161
162if __name__ == '__main__':
163    main()
164