• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
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.server import frontend
34
35
36ROLES_REQUIRE_TKO_ACCESS = {
37        'afe',
38        'database',
39        'drone',
40        'scheduler',
41        'sentinel',
42        'shard',
43}
44
45def 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
60def 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(','))
82    print 'Adding servers %s to access list for projects %s' % (servers,
83                                                                instance)
84    print 'Fetching their IP addresses...'
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
92    print '...Done: %s' % ips
93
94    cidr_ips = [str(ip) + '/32' for ip in ips]
95
96    login = False
97    while True:
98        try:
99            utils.run('gcloud config set project %s -q' % project)
100            cmd = ('gcloud sql instances patch %s --authorized-networks %s '
101                   '-q' % (instance, ','.join(cidr_ips)))
102            print 'Running command to update whitelists: "%s"' % cmd
103            utils.run(cmd, stdout_tee=sys.stdout, stderr_tee=sys.stderr)
104            return
105        except error.CmdError:
106            if login:
107                raise
108
109            # Try to login and retry if the command failed.
110            gcloud_login(project)
111            login = True
112
113
114def main():
115    """main script."""
116    parser = argparse.ArgumentParser()
117    parser.add_argument('--project', type=str, dest='project',
118                        help='Name of the Google Cloud project.')
119    parser.add_argument('--instance', type=str, dest='instance',
120                        help='Name of the CloudSQL instance.')
121    parser.add_argument('--afe', type=str, dest='afe',
122                        help='Name of the RPC server to get server list.',
123                        default=None)
124    parser.add_argument('--extra_servers', type=str, dest='extra_servers',
125                        help=('Extra servers to be included in the "Allowed '
126                              'Networks" list separated by comma.'),
127                        default=None)
128    options = parser.parse_args()
129
130    update_allowed_networks(options.project, options.instance, options.afe,
131                            options.extra_servers)
132
133
134if __name__ == '__main__':
135    main()
136