def probe_regions(regions=None, tag_filter=None):
"""probe AWS region(s) for instances, and return dict(s) of instance(s)
Requires ec2:DescribeInstances permission.
Uses multithreading to probe all whitelisted regions simultaneously.
Args:
regions (list[str]): AWS region(s) to probe.
tag_filter (list[dict]): Tag filter to filter instances with.
Returns:
list[dict]: Found instance(s).
'region' (str): AWS region that an instance is in.
For other key-value pairs, see what _probe_region returns.
"""
if regions is None:
regions = consts.REGIONS
if tag_filter is None:
tag_filter = []
tag_filter.append({'Name': "tag:Namespace", 'Values': [consts.NAMESPACE]})
threader = Threader()
for region in regions:
threader.add_thread(_probe_region, (region, tag_filter))
regions_instances = threader.get_results(return_dict=True)
return [{
'region': region,
**instance
} for region, instances in regions_instances.items()
for instance in instances]
def test_threader_fifo():
"""test that threader is first-in, first-out"""
def func(index):
sleep(randint(5, 25) / 100)
return index
threader = Threader()
for index in range(10):
threader.add_thread(func, (index, ))
assert threader.get_results() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
def delete_component(self):
"""delete VPC(s) and associated SG(s) from AWS"""
threader = Threader()
for region in consts.REGIONS:
threader.add_thread(self._delete_region_vpc, (region, ))
deleted_vpcs = threader.get_results()
if any(deleted_vpcs):
print(f"VPC {consts.NAMESPACE} deleted "
"from whitelisted AWS region(s).")
else:
print("No VPCs to delete.")
def delete_component(self):
"""remove namespace EC2 key pair from each whitelisted AWS region"""
threader = Threader()
for region in consts.REGIONS:
threader.add_thread(self._delete_region_key_pair, (region, ))
deleted_key_pairs = threader.get_results()
if any(deleted_key_pairs):
print(f"EC2 key pair {self._key_pair_name} "
"deleted from whitelisted AWS region(s).")
else:
print("No EC2 key pairs to delete.")
def check_component(self):
"""determine which regions need namespace EC2 key pair created
Returns:
dict: Which regions namespace EC2 key pair exists in.
Region name (str/None): Public key fingerprint, if pair exists.
"""
threader = Threader()
for region in consts.REGIONS:
threader.add_thread(self._region_namespace_key_fingerprint,
(region, ))
return threader.get_results(return_dict=True)
def test_threader_dict_return():
"""test that functions' first args set as function returns' dict keys"""
def func(dict_key):
return dict_key * 2
threader = Threader()
for index in range(5):
threader.add_thread(func, (index, ))
assert threader.get_results(return_dict=True) == {
0: 0,
1: 2,
2: 4,
3: 6,
4: 8
}
def probe_regions():
"""return elastic IP addresses from whitelisted regions
Requires ec2:DescribeInstances and ec2:DescribeAddresses permissions.
"""
all_instances = find_instances.probe_regions()
threader = Threader()
for region in consts.REGIONS:
threader.add_thread(_probe_region, (region, all_instances))
region_addresses = threader.get_results(return_dict=True)
return [{'region': region, **address}
for region, addresses in region_addresses.items()
for address in addresses]
def upload_component(self, fingerprint_regions):
"""create namespace EC2 key pair in each whitelisted AWS region
Args:
fingerprint_regions (dict): See what check_component returns.
"""
aws_fingerprints = [
fp for fp in fingerprint_regions.values() if fp is not None
]
if consts.RSA_KEY_PEM.is_file():
pub_key_bytes = pem.pem_to_public_key()
print(f"Using existing {self._pem_file} file for EC2 key pair(s).")
# If SSH key pair doesn't exist in any regions, create a new one
elif not aws_fingerprints:
pub_key_bytes = pem.generate_rsa_key_pair()
print(f"Generating new {self._pem_file} file for EC2 key pair(s).")
# No private key file, and there are existing EC2 key pairs
else:
halt.err(
f"RSA private key file {self._pem_file} not found.",
" Additional pairs must be created from same private key.")
if len(set(aws_fingerprints)) > 1:
print("Warning: Differing EC2 key pairs found.")
local_key_fingerprint = pem.local_key_fingerprint()
if local_key_fingerprint not in aws_fingerprints and aws_fingerprints:
halt.err("Local key fingerprint doesn't match any EC2 key pair.")
threader = Threader()
for region in fingerprint_regions:
if fingerprint_regions[region] is None:
threader.add_thread(self._create_region_key_pair,
(region, pub_key_bytes))
created_pair_fingerprints = threader.get_results()
if created_pair_fingerprints:
print(f"EC2 key pair {self._key_pair_name} created in "
f"{len(created_pair_fingerprints)} AWS region(s).")
else:
print(f"EC2 key pair {self._key_pair_name} "
"already present in whitelisted region(s).")
def main(regions: List[str]) -> str:
"""repeatedly use ec2:DescribeRegions action to estimate closest region"""
ping_num = 50
threader = Threader()
for region in regions:
ec2_client = aws.ec2_client_no_validate(region)
for _ in range(ping_num):
threader.add_thread(_get_region_latency, (ec2_client,))
latencies = threader.get_results()
latencies_for_regions = []
for i, region in enumerate(regions):
region_latencies = sorted(latencies[i*ping_num:(i+1)*ping_num])
region_latencies = region_latencies[ping_num//5:-ping_num//5]
region_latency = sum(region_latencies) / len(region_latencies)
latencies_for_regions.append((region, region_latency))
latencies_for_regions.sort(key=lambda tup: tup[1])
return latencies_for_regions[0][0]
def test_threader_add_thread_type_validation():
"""test that add_thread only accepts function and tuple as args 1 and 2"""
def func(index):
return index
with pytest.raises(ValueError) as excinfo:
threader = Threader()
threader.add_thread("I'm a string!", ("1st arg", ))
assert str(excinfo.value) == "func must be a function."
with pytest.raises(ValueError) as excinfo:
threader = Threader()
threader.add_thread(func, "fargs...")
assert str(excinfo.value) == "fargs must be a non-empty tuple."
with pytest.raises(ValueError) as excinfo:
threader = Threader()
threader.add_thread(func, tuple())
assert str(excinfo.value) == "fargs must be a non-empty tuple."
def check_component(self):
"""determine statuses of VPC(s) and SG(s) on AWS
Returns:
tuple:
dict: Which regions namespace VPC exists in.
'ToCreate' (list): AWS region(s) to create VPC in.
'Existing' (list): AWS region(s) already containing VPC.
dict: VPC security group status(es) for each region.
Name of security group (dict):
'ToCreate' (list): AWS region(s) to create SG in.
'ToUpdate' (list): AWS region(s) to update SG in.
'UpToDate' (list): AWS region(s) SG is up to date in.
"""
regions = consts.REGIONS
# Region(s) to create VPC in, and region(s) already containing VPC
vpc_regions = {'ToCreate': list(regions), 'Existing': []}
# Status for each SG in each region
sg_names = {
sg_name: {
'ToCreate': list(regions),
'ToUpdate': [],
'UpToDate': []
}
for sg_name in self._security_group_setup
}
vpc_threader = Threader()
for region in regions:
vpc_threader.add_thread(aws.get_region_vpc, (region, ))
# VPCs already present in AWS regions
aws_vpcs = vpc_threader.get_results(return_dict=True)
# Check each region has VPC with correct Name tag value
for region in regions:
if aws_vpcs[region] is not None:
if ({
'Key': "Name",
'Value': consts.NAMESPACE
} in aws_vpcs[region]['Tags']):
vpc_regions['ToCreate'].remove(region)
vpc_regions['Existing'].append(region)
# TODO: Detect and repair incomplete VPCs (missing subnets, etc.)
sg_threader = Threader()
for region in regions:
if aws_vpcs[region] is not None:
sg_threader.add_thread(aws.get_vpc_security_groups,
(region, aws_vpcs[region]['VpcId']))
# VPC security groups already present in AWS regions
aws_sgs = sg_threader.get_results(return_dict=True)
# Check each region for VPC SG(s) described by aws_setup.json
for sg_name, sg_regions in sg_names.items():
for region in regions:
if aws_vpcs[region] is not None:
if any(aws_sg['GroupName'] == sg_name
for aws_sg in aws_sgs[region]):
sg_regions['ToCreate'].remove(region)
sg_regions['ToUpdate'].append(region)
# Check if existing VPC SG needs to be updated
if region in sg_regions['ToUpdate']:
local_sg_ingress = self._get_json_sg_ingress(sg_name)
aws_sg_ingress = next(sg['IpPermissions']
for sg in aws_sgs[region]
if sg['GroupName'] == sg_name)
sg_ingress_diffs = DeepDiff(local_sg_ingress,
aws_sg_ingress,
ignore_order=True)
if not sg_ingress_diffs:
sg_regions['ToUpdate'].remove(region)
sg_regions['UpToDate'].append(region)
return (vpc_regions, sg_names)
def upload_component(self, vpc_and_sg_info):
"""create VPC(s) and create/update SG(s) in AWS region(s)
Args:
vpc_and_sg_info (dict): See what check_component returns.
"""
vpc_regions, sg_names = vpc_and_sg_info
vpc_threader = Threader()
for region in vpc_regions['ToCreate']:
vpc_threader.add_thread(self._create_vpc, (region, ))
vpc_threader.get_results()
create_num = len(vpc_regions['ToCreate'])
if create_num > 0:
print(f"VPC {consts.NAMESPACE} created in {create_num} region(s).")
else:
print(f"VPC {consts.NAMESPACE} already present "
"in whitelisted region(s).")
vpc_ids = {}
threader = Threader()
for region in consts.REGIONS:
threader.add_thread(aws.get_region_vpc, (region, ))
for region, vpc in threader.get_results(return_dict=True).items():
if vpc is None:
halt.err(f"Namespace VPC not created in {region} region.")
vpc_ids[region] = vpc['VpcId']
sg_threader = Threader()
for sg_name, sg_regions in sg_names.items():
sg_desc = self._security_group_setup[sg_name]
for region in sg_regions['ToCreate']:
sg_threader.add_thread(
self._create_sg,
(region, sg_name, sg_desc, vpc_ids[region]))
for region in sg_regions['ToUpdate']:
sg_threader.add_thread(self._update_sg,
(region, sg_name, vpc_ids[region]))
sg_threader.get_results()
for sg_name, sg_regions in sg_names.items():
if sg_regions['ToCreate']:
print(f"VPC SG {sg_name} created in "
f"{len(sg_regions['ToCreate'])} region(s).")
if sg_regions['ToUpdate']:
print(f"VPC SG {sg_name} updated in "
f"{len(sg_regions['ToUpdate'])} region(s).")
if not sg_regions['ToCreate'] and not sg_regions['ToUpdate']:
print(f"VPC SG {sg_name} already up to date "
"in whitelisted region(s).")