var AWS = require('aws-sdk');

const SG = '<%= security_group_name %>';

const ports = <%= security_group_ports %>;

const ec2 = new AWS.EC2({ region: process.env.AWS_REGION });

exports.handler = function(event, context) {
    var d = get_desired_sg();
    var c = get_current_sg();

    Promise.all([d, c]).then(returns => {
        var add_rules = [...returns[0]].filter(x => !returns[1][0].has(x));
        var remove_rules = [...returns[1][0]].filter(x => !returns[0].has(x));

        var authorize = authorize_ips(add_rules, returns[1][1]);
        var revoke = revoke_ips(remove_rules, returns[1][1]);

        Promise.all([authorize, revoke]).then(() => { context.succeed("cool"); });
    });
};


revoke_ips = (remove_rules, group_id) => {
    if (remove_rules.length === 0) return Promise.resolve();
    console.log("REVOKING:");
    console.log(remove_rules);
    return ec2.revokeSecurityGroupIngress(get_sg_params(remove_rules, group_id)).promise();
};


authorize_ips = (add_rules, group_id) => {
    if (add_rules.length === 0) return Promise.resolve();
    console.log("AUTHORIZING:");
    console.log(add_rules);
    return ec2.authorizeSecurityGroupIngress(get_sg_params(add_rules, group_id)).promise();
};


get_sg_params = (ip_ports, group_id) => {
    return {
        DryRun: false,
        GroupId: group_id,
        IpPermissions: ip_ports.map(ip_port => {
            tuple = ip_port.split(":");
            return {
                IpProtocol: 'tcp',
                FromPort: tuple[1],
                ToPort: tuple[1],
                IpRanges: [{
                    CidrIp: tuple[0]
                }]
            };
        })
    };
};


get_current_sg = () => {
    return new Promise((resolve, reject) => {
        var params = { Filters: [{ Name: 'tag:Name', Values: [SG] }] };
        var current_sg = [];

        ec2.describeSecurityGroups(params, (err, data) => {
            if (err) reject(err);
            else {
                data.SecurityGroups[0].IpPermissions.forEach((rule) => {
                    var port = rule.FromPort;
                    while (port <= rule.ToPort) {
                        rule.IpRanges.forEach((ip) => {
                            current_sg.push(ip.CidrIp + ":" + port);
                        });
                        port++;
                    }
                });
                resolve([new Set(current_sg), data.SecurityGroups[0].GroupId]);
            }
        });
    });
};


const get_desired_sg = () => {
    return new Promise((resolve, reject) => {
        var cpetools_creds = new AWS.TemporaryCredentials({
            RoleArn: 'arn:aws:iam::994136867826:role/DiscoverTcagentsIAMRole-USWest2-a23698463a8b1d4e2036d66fbd69aef'
        });

        var cpetools_ec2 = new AWS.EC2({ credentials: cpetools_creds, region: 'us-west-2' });

        var params = {
            Filters: [
                { 'Name': 'tag:aws:cloudformation:stack-name', 'Values': ['TeamcityAgentService'] },
                { 'Name': 'instance-state-name', 'Values': ['running'] }
            ]
        };

        var tcagent_ips = [];

        cpetools_ec2.describeInstances(params, (err, data) => {
            if (err) reject(err);
            else {
                data.Reservations.forEach((r) => {
                    r.Instances.forEach((i) => {
                        ports.forEach((p) => {
                            tcagent_ips.push(i.PublicIpAddress + "/32:" + p);
                        });
                    });
                });
                resolve(new Set(tcagent_ips));
            }
        });
    });
};
