#!/usr/bin/env groovy
/*
* Licensed Materials - Property of IBM Corp.
* IBM UrbanCode Build
* IBM UrbanCode Deploy
* IBM UrbanCode Release
* IBM AnthillPro
* (c) Copyright IBM Corporation 2002, 2013. All Rights Reserved.
*
* U.S. Government Users Restricted Rights - Use, duplication or disclosure restricted by
* GSA ADP Schedule Contract with IBM Corp.
*/

import com.urbancode.air.AirPluginTool
import com.urbancode.air.XTrustProvider

final def apTool = new AirPluginTool(this.args[0], this.args[1])

final def props = apTool.getStepProperties()

final def serverUrlArray = props['serverURL']?.split('://') as String[]
final def username = props['username']
final def password = props['password']
final def timeout = Integer.parseInt(props['timeout'])

final def poolNameList = props['poolName']
def poolNames = (poolNameList.tokenize('\n,').findAll {it && it.trim().length() > 0}).collect {it.trim()}
final def isPoolGlobal = props['isPoolGlobal']?.toBoolean()
final def nodeAddress = props['nodeAddress'] as String
final def nodePort = Long.parseLong(props['nodePort']) as long
final def doWait = Boolean.valueOf(props['doWaitForConnectionClose'])
final def forceOffline = props['forceOffline']?.toBoolean()
final def sleepInterval = (props['sleepInterval'])?Long.valueOf(props['sleepInterval']):5

try {
    if (serverUrlArray.length != 2)
        println("Malformed server URL: ${props['serverURL']} missing <protocol>://")
    if (serverUrlArray[1].endsWith('/'))
        serverUrlArray[1] = serverUrlArray[1].substring(0, serverUrlArray[1].length() - 1)

    final def connectionString = serverUrlArray[0] + "://" + username + ":" + password + "@" +
            serverUrlArray[1] + "/iControl/iControlPortal.cgi"

    XTrustProvider.install()
    def poolMemberAddress

    def getState = { poolName ->
        if (isPoolGlobal) {
            for (member in poolMemberAddress.get_enabled_state(
                    [poolName] as String[],
                    [new iControl.CommonIPPortDefinition(nodeAddress, nodePort)] as iControl.CommonIPPortDefinition[][])[0]) {
                if (member.member.address.equals(nodeAddress) && member.member.port == nodePort) return member.state
            }
        }
        else {
            for (member in poolMemberAddress.get_session_enabled_state(poolName)[0]) {
                if (member.member.address.equals(nodeAddress) && member.member.port == nodePort) return member.session_state
            }
        }
    }

    def getActiveConnections = { poolName ->
        for (member in poolMemberAddress.get_statistics(
                [poolName] as String[],
                [new iControl.CommonIPPortDefinition(nodeAddress, nodePort)] as iControl.CommonIPPortDefinition[][])[0]
                    .statistics[0].statistics) {
            if (member.type == iControl.CommonStatisticType.STATISTIC_SERVER_SIDE_CURRENT_CONNECTIONS)
                return member.value.low
        }
    }

    def removeNode = {
        //------------------------------------------------------------------------------
        // Remove node from pool
        //------------------------------------------------------------------------------
        println("Connecting to server ${props['serverURL']} using timeout of $timeout second(s)")

        if (isPoolGlobal) {
            poolMemberAddress =
                new iControl.GlobalLBPoolMemberLocator().getGlobalLBPoolMemberPort(new java.net.URL(connectionString)) as
                    iControl.GlobalLBPoolMemberBindingStub
        }
        else {
            poolMemberAddress =
                new iControl.LocalLBPoolMemberLocator().getLocalLBPoolMemberPort(new java.net.URL(connectionString)) as
                    iControl.LocalLBPoolMemberBindingStub
        }
        poolMemberAddress.setTimeout(timeout * 1000);
        poolNames.each { poolName ->
            if (poolName && poolName.trim().length() > 0) {
                def currentState = getState(poolName)
                println("Current node state in pool $poolName: $currentState")
                if (currentState == iControl.CommonEnabledState.STATE_ENABLED) {
                    println("Disabling node $nodeAddress:$nodePort from pool $poolName")
                    if (isPoolGlobal) {
                        poolMemberAddress.set_enabled_state(
                                [poolName] as String[],
                                [new iControl.GlobalLBPoolMemberMemberEnabledState(
                                        new iControl.CommonIPPortDefinition(nodeAddress, nodePort),
                                        iControl.CommonEnabledState.STATE_DISABLED)
                                ] as iControl.GlobalLBPoolMemberMemberEnabledState[][])
                    }
                    else {
                        poolMemberAddress.set_session_enabled_state(
                                [poolName] as String[],
                                [new iControl.LocalLBPoolMemberMemberSessionState(
                                        new iControl.CommonIPPortDefinition(nodeAddress, nodePort),
                                        iControl.CommonEnabledState.STATE_DISABLED)
                                ] as iControl.LocalLBPoolMemberMemberSessionState[][])
                        if (forceOffline) {
                            poolMemberAddress.set_monitor_state(
                                [poolName] as String[],
                                [new iControl.LocalLBPoolMemberMemberMonitorState(
                                        new iControl.CommonIPPortDefinition(nodeAddress, nodePort),
                                        iControl.CommonEnabledState.STATE_DISABLED)
                                ] as iControl.LocalLBPoolMemberMemberMonitorState[][])
                            println "Node forced offline."
                        }

                    }
                    println("New node state in pool $poolName: ${getState(poolName)}")
                }
                else if (currentState == null) {
                    println("Node $nodeAddress:$nodePort was not found in pool $poolName")
                }
                else {
                    println("Node $nodeAddress:$nodePort is not enabled for pool $poolName so nothing to do")
                }
            }
        }
    }

    def waitForConnections = { poolName ->
        //------------------------------------------------------------------------------
        // Wait for connections to close
        //------------------------------------------------------------------------------
        def currentConnections = getActiveConnections(poolName)
        while(currentConnections > 0) {
            println("Waiting for $currentConnections to close, sleep for $sleepInterval second(s)")
            sleep(sleepInterval * 1000, {println("Waiting aborted!"); return true})
            currentConnections = getActiveConnections(poolName)
        }
        println("All connections closed for pool $poolName, done waiting.")
    }

    removeNode()
    if (doWait) {
        poolNames.each { poolName ->
            println "Waiting for connections in pool $poolName"
            waitForConnections(poolName)
        }
    }
}
catch (Exception e) {
    e.printStackTrace()
    System.exit(1)
}
