github.com/rochacon/deis@v1.0.2-0.20150903015341-6839b592a1ff/contrib/linode/create-linode-user-data.py (about)

     1  #!/usr/bin/env python
     2  """
     3  Create CoreOS user-data by merging contrib/coreos/user-data and contrib/linode/linode-user-data-template
     4  
     5  Usage: create-linode-user-data.py
     6  """
     7  import base64
     8  import sys
     9  import os
    10  import collections
    11  import argparse
    12  
    13  import yaml
    14  import colorama
    15  from colorama import Fore, Style
    16  import requests
    17  
    18  
    19  def combine_dicts(orig_dict, new_dict):
    20      for key, val in new_dict.iteritems():
    21          if isinstance(val, collections.Mapping):
    22              tmp = combine_dicts(orig_dict.get(key, {}), val)
    23              orig_dict[key] = tmp
    24          elif isinstance(val, list):
    25              orig_dict[key] = (orig_dict.get(key, []) + val)
    26          else:
    27              orig_dict[key] = new_dict[key]
    28      return orig_dict
    29  
    30  
    31  def get_file(name, mode="r", abspath=False):
    32      current_dir = os.path.dirname(__file__)
    33  
    34      if abspath:
    35          return file(os.path.abspath(os.path.join(current_dir, name)), mode)
    36      else:
    37          return file(os.path.join(current_dir, name), mode)
    38  
    39  
    40  def main():
    41      colorama.init()
    42  
    43      parser = argparse.ArgumentParser(description='Create Linode User Data')
    44      parser.add_argument('--public-key', action='append', required=True, type=file, dest='public_key_files', help='Authorized SSH Keys')
    45      parser.add_argument('--etcd-token', required=False, default=None, dest='etcd_token', help='Etcd Token')
    46      args = parser.parse_args()
    47  
    48      etcd_token = args.etcd_token
    49      if etcd_token is None:
    50          etcd_token = generate_etcd_token()
    51      else:
    52          if not validate_etcd_token(args.etcd_token):
    53              raise ValueError('Invalid Etcd Token. You can generate a new token at https://discovery.etcd.io/new.')
    54  
    55      public_keys = []
    56      for public_key_file in args.public_key_files:
    57          public_key = public_key_file.read()
    58          if validate_public_key(public_key):
    59              public_keys.append(public_key)
    60          else:
    61              log_warning('Invalid public key: ' + public_key_file.name)
    62  
    63      if not len(public_keys) > 0:
    64          raise ValueError('Must supply at least one valid public key')
    65  
    66      linode_user_data = get_file("linode-user-data.yaml", "w", True)
    67      linode_template = get_file("linode-user-data-template.yaml")
    68      coreos_template = get_file("../coreos/user-data.example")
    69  
    70      configuration_linode_template = yaml.safe_load(linode_template)
    71      configuration_coreos_template = yaml.safe_load(coreos_template)
    72  
    73      configuration = combine_dicts(configuration_coreos_template, configuration_linode_template)
    74      configuration['ssh_authorized_keys'] = public_keys
    75  
    76      dump = yaml.dump(configuration, default_flow_style=False, default_style='|')
    77      dump = dump.replace('#DISCOVERY_URL', 'https://discovery.etcd.io/' + str(etcd_token))
    78  
    79      with linode_user_data as outfile:
    80          outfile.write("#cloud-config\n\n" + dump)
    81          log_success('Wrote Linode user data to ' + linode_user_data.name)
    82  
    83  
    84  def validate_public_key(key):
    85      try:
    86          type, key_string, comment = key.split()
    87          data = base64.decodestring(key_string)
    88          return data[4:11] == type
    89      except:
    90          return False
    91  
    92  
    93  def generate_etcd_token():
    94      log_info('Generating new Etcd token...')
    95      data = requests.get('https://discovery.etcd.io/new').text
    96      token = data.replace('https://discovery.etcd.io/', '')
    97      log_success('Generated new token: ' + token)
    98      return token
    99  
   100  
   101  def validate_etcd_token(token):
   102      try:
   103          int(token, 16)
   104          return True
   105      except:
   106          return False
   107  
   108  
   109  def log_info(message):
   110      print(Fore.CYAN + message + Fore.RESET)
   111  
   112  
   113  def log_warning(message):
   114      print(Fore.YELLOW + message + Fore.RESET)
   115  
   116  
   117  def log_success(message):
   118      print(Style.BRIGHT + Fore.GREEN + message + Fore.RESET + Style.RESET_ALL)
   119  
   120  
   121  def log_error(message):
   122      print(Style.BRIGHT + Fore.RED + message + Fore.RESET + Style.RESET_ALL)
   123  
   124  
   125  if __name__ == "__main__":
   126      try:
   127          main()
   128      except Exception as e:
   129          log_error(e.message)
   130          sys.exit(1)