'use strict';

const AWS = require('aws-sdk');
const https = require('https');

const CONFIG = {
    groupId: process.env.SECURITY_GROUP_ID || '',
    protocol: process.env.PROTOCOL || 'tcp',
    port: process.env.PORT || 80,
    s3bucket: process.env.S3BUCKET || ''
};

const EC2 = new AWS.EC2();
const S3 = new AWS.S3();

/**
 * Promiseのログを残す共通関数
 *
 * @param {Promise} promise
 */
function log(promise) {
    return promise.then(data => {
        // 正常時は取得したデータをログに出力
        console.log(data);

        return data;
    }).catch(error => {
        // エラー時はエラーの内容をログに出力
        console.log(error, error.stack);

        return error;
    });
}

/**
 * セキュリティグループを操作する共通関数
 *
 * @param {String} method - 動作 'authorizeSecurityGroupIngress' or 'revokeSecurityGroupIngress'
 * @param {Array} ipRanges - IPアドレス範囲の配列 ex. ['192.168.0.0/24','8.8.8.8/32','8.8.4.4/32']
 * @return {Promise}
 */
function securityGroup(method, ipRanges) {

    // IPアドレス範囲の配列をオプションとして渡せる形式に変換
    let optionsIpRanges = [];
    ipRanges.forEach(ipRange => {
        optionsIpRanges.push({CidrIp: ipRange});
    });

    // セキュリティグループを操作するAWS-SDKのメソッドを実行
    return log(EC2[method]({
        GroupId: CONFIG.groupId,
        IpPermissions: [{
            IpProtocol: CONFIG.protocol,
            FromPort: CONFIG.port,
            ToPort: CONFIG.port,
            IpRanges: optionsIpRanges
        }]
    }).promise());
}

/**
 * S3を操作する共通関数
 *
 * @param {String} method - 動作 'putObject' or 'getObject'
 * @param {String} key - セット or ゲットするキー
 * @param {String} [body] - セットする内容（ゲットする場合は不要）
 * @return {Promise}
 */
function s3(method, key, body) {
    let options = {
        Bucket: CONFIG.s3bucket,
        Key: key
    };
    if (method === 'putObject') options.Body = body;

    // S3を操作するAWS-SDKのメソッドを実行
    return log(S3[method](options).promise().then(data => {
        return (method === 'getObject') ? data.Body.toString() : data
    }));
}

/**
 * S3にIP範囲の情報を追加する関数
 *
 * @param {String} syncToken - 追加するip-range.jsonのsyncToken
 * @param {Array} body - 追加するIP範囲の配列
 * @return {Promise}
 */
function addIpRangeToS3(syncToken, body) {
    return s3('putObject', syncToken, JSON.stringify(body));
}

/**
 * S3上のHEADファイル（現在のip-range.jsの位置を示すファイル）を設定・更新する関数
 *
 * @param {String} syncToken - セットする新しいsyncToken
 * @return {Promise}
 */
function setTheHeadFileOnS3(syncToken) {
    return s3('putObject', 'HEAD', syncToken);
}

/**
 * S3上のHEADファイル（現在のip-range.jsの位置を示すファイル）を取得する関数
 *
 * @return {Promise}
 */
function getTheHeadFileOnS3() {
    return s3('getObject', 'HEAD');
}

/**
 * S3のIP範囲の情報を取得する関数
 *
 * @param {String} syncToken - 取得するIP範囲のsyncToken
 * @return {Promise}
 */
function getIpRangeFromS3(syncToken) {
    return s3('getObject', syncToken).then(data => {
        return JSON.parse(data);
    });
}

/**
 * JSONファイルを取得する関数（httpsのみ）
 *
 * @param {String} url - 取得するjsonファイルのURL
 * @return {Promise}
 */
function getJSON(url) {
    return new Promise((resolve, reject) => {
        https.get(url, res => {
            let body = '';
            res.setEncoding('utf8');
            res.on('data', chunk => {
                body += chunk;
            });
            res.on('end', () => {
                resolve(JSON.parse(body));
            });
        }).on('error', e => {
            reject(e);
        });
    });
}

/**
 * ip-range.jsonファイルをCloudFrontのIP範囲の配列に変換する関数
 *
 * @param {Object} json - ip-range.jsonの内容
 * @return {Array} ipRange - CloudFrontのIP範囲の配列
 */
function changeIpRangeJsonToCloudFrontIpRangeArray(json) {
    let ipRange = [];
    for (let i = 0; i < json.prefixes.length; i++) {
        if (json.prefixes[i].service === "CLOUDFRONT") {
            ipRange.push(json.prefixes[i].ip_prefix);
        }
    }
    return ipRange;
}

/**
 * 古いIP範囲と新しいIP範囲の配列からセキュリティグループの内容を変更する関数
 *
 * @param {Array} oldIpRange - 古いIP範囲の配列 ex. ['192.168.0.0/24','8.8.8.8/32','8.8.4.4/32']
 * @param {Array} newIpRange - 新しいIP範囲の配列 ex. ['192.168.0.0/24','8.8.8.8/32','8.8.4.4/32']
 * @return {Promise}
 */
function changeIpRangeOfSecurityGroup(oldIpRange, newIpRange) {

    // 差分のみ追加・削除を実行するように、配列の重複を削除
    let i = 0;
    while (i < oldIpRange.length) {
        let j = 0;
        while (j < newIpRange.length) {
            if (oldIpRange[i] == newIpRange[j]) {
                oldIpRange.splice(i, 1);
                newIpRange.splice(j, 1);
                i--;
                break;
            }
            j++;
        }
        i++;
    }

    // セキュリティグループの操作
    let promises = [];
    if (oldIpRange.length > 0) promises.push(securityGroup('revokeSecurityGroupIngress', oldIpRange));
    if (newIpRange.length > 0) promises.push(securityGroup('authorizeSecurityGroupIngress', newIpRange));
    return Promise.all(promises);
}


/********************************************************************************************************************
 * 以下、Lambda実行時のメイン関数                                                                                       *
 ********************************************************************************************************************/
exports.handler = function (event, context, callback) {
    /**
     * AmazonIpSpaceChangedトピックから通知を受け取ると、event.Records[0].Sns.Messageには、以下の情報が文字列で格納される
     * {
     *   "create-time":"yyyy-mm-ddThh:mm:ss+00:00",
     *   "synctoken":"0123456789",
     *   "md5":"6a45316e8bc9463c9e926d5d37836d33",
     *   "url":"https://ip-ranges.amazonaws.com/ip-ranges.json"
     * }
     */
    const message = JSON.parse(event.Records[0].Sns.Message);

    if (typeof message != 'object' || message == null) return;

    /**
     * 共通処理をまとめた関数
     *
     * @param {Array} oldIpRange - 古い（削除する）IP範囲の配列 ex. ['192.168.0.0/24','8.8.8.8/32','8.8.4.4/32']
     * @return {Promise}
     */
    function process(oldIpRange) {
        // ip-range.jsonを取得
        return getJSON(message.url).then(json => {
            // 新しく登録するIP範囲の配列を取得
            let newIpRange = changeIpRangeJsonToCloudFrontIpRangeArray(json);

            return Promise.all([
                // syncTokenをS3のHEADファイルに保存
                setTheHeadFileOnS3(json.syncToken),

                // 新しく登録するIP範囲の情報をS3に保存
                addIpRangeToS3(json.syncToken, newIpRange),

                // セキュリティグループの情報を変更
                changeIpRangeOfSecurityGroup(oldIpRange, newIpRange)
            ]).then(() => {
                callback(null, 'success')
            }).catch(error => {
                callback(Error, error);
            });
        });
    }

    // 現在設定されているセキュリティグループの情報（syncToken）を取得
    getTheHeadFileOnS3().then(head => {
        if (head !== message.synctoken) {
            // 古い（削除する）IP範囲を取得
            getIpRangeFromS3(head).then(oldIpRange => {
                process(oldIpRange);
            });
        }
    }).catch(() => {
        // エラー時（HEADファイルがないときなど）

        // 削除するIP範囲はなしで、共通関数実行
        process([]);
    });
};