import boto3 import datetime import json import sys import time REGION_NAME = 'region_name' REGION_KEY = 'region_key' REGION_SECURITY_GROUP = 'region_security_group' REGION_HUMAN_NAME = 'region_human_name' INSTANCE_TYPE = 't2.micro' REGION_AMI = 'region_ami' def read_region_config(region_config='configuration.txt'): config = {} with open(region_config, 'r') as f: for myline in f: mylist = [item.strip() for item in myline.strip().split(',')] region_num = mylist[0] config[region_num] = {} config[region_num][REGION_NAME] = mylist[1] config[region_num][REGION_KEY] = mylist[2] config[region_num][REGION_SECURITY_GROUP] = mylist[3] config[region_num][REGION_HUMAN_NAME] = mylist[4] config[region_num][REGION_AMI] = mylist[5] return config def get_ip_list(response): if response.get('Instances', None): return [instance.get('PublicIpAddress', None) for instance in response['Instances']] else: return [] def create_ec2_client(region_number, region_config): config = read_region_config(region_config) region_name = config[region_number][REGION_NAME] session = boto3.Session(region_name=region_name) return session.client('ec2') def collect_public_ips(region_number, node_name_tag, region_config): ec2_client = create_ec2_client(region_number, region_config) filters = [{'Name': 'tag:Name','Values': [node_name_tag]}] response = ec2_client.describe_instances(Filters=filters) ip_list = [] for reservation in response[u'Reservations']: ip_list.extend(instance['PublicIpAddress'] for instance in reservation['Instances']) return ip_list def generate_distribution_config2(region_number, node_name_tag, region_config, shard_number, client_number, distribution_config): ip_list = collect_public_ips(region_number, node_name_tag, region_config) generate_distribution_config(shard_number, client_number, ip_list, distribution_config) def generate_distribution_config3(shard_number, client_number, ip_list_file, distribution_config): with open(ip_list_file, "r") as fin: lines = fin.readlines() ip_list = [line.strip() for line in lines] generate_distribution_config(shard_number, client_number, ip_list, distribution_config) def generate_distribution_config(shard_number, client_number, ip_list, distribution_config): if len(ip_list) < shard_number * 2 + client_number: print("Not enough nodes to generate a config file") return False # Create ip for clients. client_id, leader_id, validator_id = 0, 0, 0 with open(distribution_config, "w") as fout: for i in range(len(ip_list)): if client_id < client_number: fout.write("%s 9000 client %d\n" % (ip_list[i], client_id % shard_number)) client_id = client_id + 1 elif leader_id < shard_number: fout.write("%s 9000 leader %d\n" % (ip_list[i], leader_id)) leader_id = leader_id + 1 else: fout.write("%s 9000 validator %d\n" % (ip_list[i], validator_id % shard_number)) validator_id = validator_id + 1 def get_availability_zones(ec2_client): response = ec2_client.describe_availability_zones() all_zones = [] if response.get('AvailabilityZones', None): all_zones = [info['ZoneName'] for info in response.get('AvailabilityZones') if info['State'] == 'available'] return all_zones def get_one_availability_zone(ec2_client): all_zones = get_availability_zones(ec2_client) if len(all_zones) > 0: return all_zones[0] else: return None # used for testing only. # if __name__ == "__main__": # ip_list = collect_public_ips('4', "4-NODE-23-36-01-2018-07-05", "configuration.txt") # print ip_list # generate_distribution_config(2, 1, ip_list, "config_test.txt")