#!/usr/bin/python
#
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import sys, subprocess
import time

is_master = subprocess.check_output(['cat /emr/instance-controller/lib/info/instance.json | jq .isMaster'], shell=True).strip()
cluster_id = subprocess.check_output(['cat /mnt/var/lib/info/job-flow.json | jq -r .jobFlowId'], shell=True).strip()
core_count = subprocess.check_output(['aws emr describe-cluster --cluster-id %s | jq -r \'.[].InstanceGroups[] | select( .InstanceGroupType == "CORE") | .RequestedInstanceCount\'' % cluster_id], shell=True).strip()
instance_group_id = subprocess.check_output(['cat /mnt/var/lib/info/instance.json | jq -r .instanceGroupId'], shell=True).strip()
instance_group_type = subprocess.check_output(['cat /mnt/var/lib/info/job-flow.json | jq -r \'.instanceGroups | .[] | select( .instanceGroupId == \"%s\") | .instanceRole\' | tr a-z A-Z' % instance_group_id], shell=True).strip()
instance_id = subprocess.check_output(['/usr/bin/curl -s http://169.254.169.254/latest/meta-data/instance-id'], shell=True)
interface_id = subprocess.check_output(['aws ec2 describe-instances --instance-ids %s | jq .Reservations[].Instances[].NetworkInterfaces[].NetworkInterfaceId' % instance_id], shell=True).strip().strip('"')
current_tag_name = subprocess.check_output(['aws ec2 describe-tags --filters Name=resource-id,Values=%s | jq -r \'.Tags | .[] | select( .Key == \"Name\") | .Value\'' % instance_id], shell=True).strip().strip('"')

#Create Tags for EMR Master/Core Instances
subprocess.check_call(['aws ec2 create-tags --resources %s --tags Key=Name,Value=%s-%s' % (instance_id, current_tag_name, instance_group_type)], shell=True)


# sys.argv[1] = master node ip, sys.argv[2~n] = core node ip
if len(sys.argv) != int(core_count)+2:
    print "Insufficient arguments"
    sys.exit(2)

if is_master == "true":
    print "This is the MASTER node"
    private_ip = str(sys.argv[1])
    
    #Assign private IP to the MASTER instance:
    subprocess.check_call(['aws ec2 assign-private-ip-addresses --network-interface-id %s --private-ip-addresses %s' % (interface_id, private_ip)], shell=True)
    subnet_id = subprocess.check_output(['aws ec2 describe-instances --instance-ids %s | jq .Reservations[].Instances[].NetworkInterfaces[].SubnetId' % instance_id], shell=True).strip().strip('"').strip().strip('"')
    subnet_cidr = subprocess.check_output(['aws ec2 describe-subnets --subnet-ids %s | jq .Subnets[].CidrBlock' % subnet_id], shell=True).strip().strip('"')
    cidr_prefix = subnet_cidr.split("/")[1]

    #Get CORE Instance ID in parameter store
    PARAMETERS = subprocess.check_output(['aws ssm describe-parameters --parameter-filters "Key=tag:Resource,Values=CORES" --query "Parameters[*]"|jq -r .[].Name'], shell=True).split()
    #Wait till all CORE instances get registered in parameter store. 
    while len(PARAMETERS) < int(core_count):
        time.sleep(1)
        PARAMETERS = subprocess.check_output(['aws ssm describe-parameters --parameter-filters "Key=tag:Resource,Values=CORES" --query "Parameters[*]"|jq -r .[].Name'], shell=True).split()
        print('Core nodes = %s'%len(PARAMETERS))

    # use for loop put all CORE node's IP parameters
    for i in range(len(PARAMETERS)):
        subprocess.check_call(['aws ssm put-parameter --name %s --value %s --type String --overwrite' % (PARAMETERS[i],sys.argv[i+2])], shell=True);

#Add the private IP address to the default network interface:
    subprocess.check_call(['sudo ip addr add dev eth0 %s/%s' % (private_ip, cidr_prefix)], shell=True)

    #Configure iptables rules such that traffic is redirected from/to the secondary to/from the primary IP address:
    primary_ip = subprocess.check_output(['/sbin/ifconfig eth0 | grep \'inet \' | cut -d: -f2 | awk \'{ print $2}\''], shell=True).strip()
    subprocess.check_call(['sudo iptables -t nat -A PREROUTING -d %s -j DNAT --to-destination %s' % (private_ip, primary_ip)], shell=True)
    subprocess.check_call(['sudo iptables -t nat -A POSTROUTING -s %s -j SNAT --to %s' % (primary_ip, private_ip)], shell=True)

else:
    print "This is the CORE node"
    
    #Add system parameter store for ip of CORE nodes
    subprocess.check_call(['aws ssm put-parameter --name %s --value "blank" --type String --tags Key=Resource,Value=CORES' % instance_id], shell=True)
    private_ip = subprocess.check_output(['aws ssm get-parameters --names %s --query "Parameters[*]"|jq -r .[].Value' % instance_id], shell=True).strip().strip('"')
    print('Private IP = %s'%private_ip)

    #Wait till MASTER is going to register all IPs of CORE nodes in parameter store
    while True:
        if private_ip != 'blank':
            break;
        else:
            time.sleep(1)
            private_ip = subprocess.check_output(['aws ssm get-parameters --names %s --query "Parameters[*]"|jq -r .[].Value' % instance_id], shell=True).strip().strip('"')
            print('Private IP = %s'%private_ip)

    #Assign private IP to the CORE instance:
    subprocess.check_call(['aws ec2 assign-private-ip-addresses --network-interface-id %s --private-ip-addresses %s' % (interface_id, private_ip)], shell=True)
    subnet_id = subprocess.check_output(['aws ec2 describe-instances --instance-ids %s | jq .Reservations[].Instances[].NetworkInterfaces[].SubnetId' % instance_id], shell=True).strip().strip('"').strip().strip('"')
    subnet_cidr = subprocess.check_output(['aws ec2 describe-subnets --subnet-ids %s | jq .Subnets[].CidrBlock' % subnet_id], shell=True).strip().strip('"')
    cidr_prefix = subnet_cidr.split("/")[1]

    #Add the private IP address to the default network interface:
    subprocess.check_call(['sudo ip addr add dev eth0 %s/%s' % (private_ip, cidr_prefix)], shell=True)

    #Configure iptables rules such that traffic is redirected from/to the secondary to/from the primary IP address:
    primary_ip = subprocess.check_output(['/sbin/ifconfig eth0 | grep \'inet \' | cut -d: -f2 | awk \'{ print $2}\''], shell=True).strip()
    subprocess.check_call(['sudo iptables -t nat -A PREROUTING -d %s -j DNAT --to-destination %s' % (private_ip, primary_ip)], shell=True)
    subprocess.check_call(['sudo iptables -t nat -A POSTROUTING -s %s -j SNAT --to %s' % (primary_ip, private_ip)], shell=True)