Home | History | Annotate | Download | only in site_utils
      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.client.common_lib.cros import retry
     34 from autotest_lib.server import frontend
     35 
     36 
     37 ROLES_REQUIRE_TKO_ACCESS = {
     38         'afe',
     39         'database',
     40         'drone',
     41         'scheduler',
     42         'sentinel',
     43         'shard',
     44         'skylab_drone',
     45 }
     46 
     47 
     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 
     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):
     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.
     78     @param dryrun: Boolean indicating whether this is a dryrun.
     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(','))
     92     print 'Adding servers %s to access list for projects %s' % (servers,
     93                                                                 instance)
     94     print 'Fetching their IP addresses...'
     95     ips = []
     96     for name in servers:
     97         try:
     98             # collect internal ips
     99             ips.append(socket.gethostbyname(name))
    100             # collect external ips
    101             ips.append(_fetch_external_ip(name))
    102         except socket.gaierror:
    103             print 'Failed to resolve internal IP address for name %s' % name
    104             raise
    105         except error.TimeoutException:
    106             print 'Failed to resolve external IP address for %s' % name
    107             raise
    108 
    109     print '...Done: %s' % ips
    110 
    111     cidr_ips = [str(ip) + '/32' for ip in ips]
    112 
    113     if dryrun:
    114         print 'This is a dryrun: skip updating glcoud sql whitelists.'
    115         return
    116 
    117     login = False
    118     while True:
    119         try:
    120             utils.run('gcloud config set project %s -q' % project)
    121             cmd = ('gcloud sql instances patch %s --authorized-networks %s '
    122                    '-q' % (instance, ','.join(cidr_ips)))
    123             print 'Running command to update whitelists: "%s"' % cmd
    124             utils.run(cmd, stdout_tee=sys.stdout, stderr_tee=sys.stderr)
    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)
    149     parser.add_argument('--dryrun', dest='dryrun', action='store_true',
    150                         default=False,
    151                         help='Fetch IPs without updating whitelists in gcloud.')
    152     options = parser.parse_args()
    153 
    154     update_allowed_networks(options.project, options.instance, options.afe,
    155                             options.extra_servers, options.dryrun)
    156 
    157 
    158 if __name__ == '__main__':
    159     main()
    160