#!/usr/bin/env python
#
# Control access to the staging load balancer via ip_access.yml

import os
import sys
import pyrax
import yaml

# Validate ye argv

if len(sys.argv) < 2:
    raise RuntimeError("%s requires one argument: the load balancer name." %
                       sys.argv[0])
the_lbname = sys.argv[1]

# Get our bearings within the filesystem.

root = os.path.realpath(os.path.join(os.path.dirname(__file__), '..'))

# Authenticate to pyrax.

pyrax.set_setting("identity_type", "rackspace")
pyrax.set_setting("region", os.environ['RAX_REGION'])
pyrax.set_credential_file(os.path.expanduser("~/.rackspace_cloud_credentials"))

clb = pyrax.cloud_loadbalancers

# Locate the staging load balancer

the_balancer = None
for each in clb.list():
    if each.name == the_lbname:
        the_balancer = each

if not the_balancer:
    raise RuntimeError("I couldn't find the staging load balancer!")

# Load the IP whitelist.

whitelist_path = os.path.join(root, 'deploy', 'ip_access.yml')
whitelist = None

with open(whitelist_path) as inf:
    whitelist = yaml.load(inf.read())['ip_accesses']

# Add source IPs for monitoring zones.

vars_path = os.path.join(root, 'deploy', 'vars.yml')
with open(vars_path) as inf:
    zone_ids = yaml.load(inf.read())['monitoring_zones']

cm = pyrax.cloud_monitoring
zones = cm.list_monitoring_zones()

for zone in zones:
    if zone.id in zone_ids:
        whitelist.extend(zone.source_ips)

# Create the access list.

access_list = [dict(address=addr, type='ALLOW') for addr in set(whitelist)]
access_list.append(dict(address='0.0.0.0/0', type='DENY'))

pyrax.utils.wait_until(the_balancer, "status", "ACTIVE", interval=1, attempts=30)
the_balancer.delete_access_list()

pyrax.utils.wait_until(the_balancer, "status", "ACTIVE", interval=1, attempts=30)
the_balancer.add_access_list(access_list)
