Az.NCache

1.0.15-beta1

framework/Resources/Scripts/stop_cache.ps1

Param(
    [Parameter(Mandatory = $true)]
    [string]$EnvironmentName,

    [string]$Name,
    [switch]$IsGracefulStop,
    [switch]$NoLogo,
    [int]$Port,
    [string]$Servers,
    [switch]$IsStopForMaintenance,
    [int]$Timeout,
    [PSCredential]$Credentials
)


function Get-EnvName {
    param (
        [Parameter(Mandatory = $true)]
        [string]$EnvironmentName
    )
    # Get all RGs with EnvironmentName tag
    $all_resource_groups =  Get-AzResourceGroup -ErrorAction Stop | Where-Object { $_.Tags -and $_.Tags.ContainsKey("EnvironmentName") } 

    if (-not $all_resource_groups) {
        throw "No resource groups found with tag 'EnvironmentName'."
    }

    # Match the exact EnvironmentName
    $matched_group = $all_resource_groups | Where-Object {
        $_.Tags["EnvironmentName"] -eq $EnvironmentName
    }

    if (-not $matched_group) {
        throw "No resource group found with EnvironmentName = '$EnvironmentName'."
    }

    return $matched_group.ResourceGroupName
}

function Get-VmsWithCache{
    Param(
        [Parameter(Mandatory = $true)]
        [string]$EnvironmentName,

        [Parameter(Mandatory = $true)]
        [string]$Name
    )

    #Filtering VMs using ResourceGroups and Cache Tags
    $vms = Get-AzVM -Status -ResourceGroupName $EnvironmentName -ErrorAction Stop | Where-Object { $_.Tags.ContainsKey("Caches") }
    
    #Assign Cache Names Array to Vms
    $vms = foreach ($vm in $vms){
        $vm | Add-Member -NotePropertyName "CachesList" -NotePropertyValue ($vm.Tags["Caches"] -split "," | ForEach-Object { $_.Trim() }) -Force

        $vm
    }

    if($Name){
        # Split comma-separated list into array, trim spaces
        $requestedCaches = $Name.Split(",") | ForEach-Object { $_.Trim() } | Where-Object { $_ }

        $vms = foreach ($vm in $vms) {
            # Find intersection between requested and actual
            $matched = $vm.CachesList | Where-Object { $requestedCaches -contains $_ }
            if ($matched) {
                $vm.CachesList = $matched
                $vm
            }
        }
    }

    return $vms
}

function Get-VMPrivateIps {
    param(
        [Parameter(Mandatory = $true)]
        [array]$VMs,

        [Parameter(Mandatory = $true)]
        [string]$ResourceGroupName,

        [string]$Servers
    )


    $VMPrivateIps = foreach ($vm in $VMs) {
        foreach ($nicRef in $vm.NetworkProfile.NetworkInterfaces) {
            $nicName = ($nicRef.Id -split "/")[-1]
            $nic = Get-AzNetworkInterface -Name $nicName -ResourceGroupName $ResourceGroupName -ErrorAction Stop 

            [PSCustomObject]@{
                VMName     = $vm.Name
                NicName    = $nic.Name
                CachesList     = $vm.CachesList
                PrivateIps = ($nic.IpConfigurations | Select-Object -ExpandProperty PrivateIpAddress)
                OsType = $vm.StorageProfile.OsDisk.OsType
            }
        }
    }

    # If Servers is provided, validate and filter further
    # Servers should be subset of VMs, If VMs doesn't contain a name mentioned in servers error will be thrown.
    if ($Servers) {
        $serverList = $Servers -split "," | ForEach-Object { $_.Trim() }
        $privateIps = $VMPrivateIps.PrivateIps

        $missing = $serverList | Where-Object { $_ -notin $privateIps }
        if ($missing.Count -gt 0) {
            throw "One of the server doesn't contain NCache"
        }

        $VMPrivateIps = $VMPrivateIps | Where-Object { $_.PrivateIps -in $serverList }
    }


    return $VMPrivateIps
}

function Get-CacheGrouping {
    param(
        [Parameter(Mandatory = $true)]
        [array]$VMInfo
    )

    $caches = foreach ($vm in $VMInfo) {
        foreach ($cache in $vm.CachesList) {
            [PSCustomObject]@{
                CacheName  = $cache
                VMName     = $vm.VMName
                PrivateIps = $vm.PrivateIps
                OsType     = $vm.OsType
            }
        }
    }

    $result = $caches | Group-Object -Property CacheName | ForEach-Object {
        [PSCustomObject]@{
            CacheName  = $_.Name
            FirstVM    = ($_.Group | Select-Object -First 1 -ExpandProperty VMName)
            OsType     = ($_.Group | Select-Object -First 1 -ExpandProperty OsType)
            PrivateIps = ($_.Group | Select-Object -ExpandProperty PrivateIps -Unique)
        }
    }

    return $result
}

function Get-LinuxCommand {
    param(
        [Parameter(Mandatory = $true)]
        $cacheInfo
    )

    $ipList = ($cacheInfo.PrivateIps -join ",")
    $command = "/opt/ncache/bin/tools/stop-cache $($cacheInfo.CacheName) -server `"$ipList`""

    if ($IsGracefulStop)     { $command += " -isgracefulstop" }
    if ($NoLogo)             { $command += " -nologo" }
    if ($Port)               { $command += " -port $Port" }
    if ($IsStopForMaintenance) { $command += " -isstopformaintenance" }
    if ($Timeout)            { $command += " -timeout $Timeout" }
    if ($Credentials)        { 
        $command += " -userid $($Credentials.UserName)"
        # Same password note as Windows
    }

    return $command
}

function Get-WindowsCommand {
    param(
        [Parameter(Mandatory = $true)]
        $cacheInfo
    )

    $ipList = ($cacheInfo.PrivateIps -join ",")
    $command = "Stop-Cache -name $($cacheInfo.CacheName) -server `"$ipList`""

    if ($IsGracefulStop)     { $command += " -isgracefulstop" }
    if ($NoLogo)             { $command += " -nologo" }
    if ($Port)               { $command += " -port $Port" }
    if ($IsStopForMaintenance) { $command += " -isstopformaintenance" }
    if ($Timeout)            { $command += " -timeout $Timeout" }
    if ($Credentials)        { 
        $command += " -userid $($Credentials.UserName)"
        # ⚠️ Normally you wouldn’t append the password for security
        # but if NCache supports it, you could theoretically do:
        # $command += " -Password $($Credentials.GetNetworkCredential().Password)"
    }

    return $command
}

function Get-CommandForCache {
    param (
        [Parameter(Mandatory=$true)]
        $cacheInfo
    )

    if ($cacheInfo.OsType -eq "Windows") {
        return @{
            CommandId = "RunPowerShellScript"
            VMName = $cacheInfo.FirstVM
            Script    = Get-WindowsCommand -cacheInfo $cacheInfo
        }
    }
    elseif ($cacheInfo.OsType -eq "Linux") {
        return @{
            CommandId = "RunShellScript"
            VMName = $cacheInfo.FirstVM
            Script    = Get-LinuxCommand -cacheInfo $cacheInfo
        }
    }
    else {
        throw "Unknown OS type for VM: $($Vm.Name)"
    }
}

function Get-Commands {
    Param(
        [Parameter(Mandatory=$true)]
        $cacheGroups
    )

    $commands = foreach($cacheInfo in $cacheGroups){
        Get-CommandForCache -cacheInfo $cacheInfo
    }

    return $commands
}

function Invoke-CacheCommand {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ResourceGroupName,

        [Parameter(Mandatory = $true)]
        [object]$Command
    )

    Write-Output "Executing on VM: $($Command.VMName) => $($Command.Script)"

    $result = Invoke-AzVMRunCommand `
        -ResourceGroupName $ResourceGroupName `
        -VMName $Command.VMName `
        -CommandId $Command.CommandId `
        -ScriptString $Command.Script

    # Aggregate all Messages into a single string
    $messages = $result.Value | ForEach-Object { $_.Message }
    $joined   = ($messages -join "`n").Trim()

    # Detect error from result codes
    $hasError = $false
    if ($Command.CommandId -eq "RunPowerShellScript" -and $result.Value | Where-Object { $_.Code -match "StdErr" -and -not [string]::IsNullOrWhiteSpace($_.Message) }) {
        $hasError = $true
    }
    elseif($Command.CommandId -eq "RunShellScript" ) {
        $message = $result.Value[0].Message
        
        if($message -match "Error:" -or $message -match "error:"){
            $hasError = $true
        }
        
        $parts   = $message -split "\[stderr\]"

        # If there's something after [stderr], treat it as error$res
        
        if ($parts.Count -gt 1 -and $parts[1].Trim().Length -gt 0) {
            $hasError = $true
        }
    }

    # Build and return structured PSObject
    $status = if ($hasError) { "error" } else { "success" }

    [PSCustomObject]@{
        status  = $status
        message = $joined
    }
}


function Invoke-CacheCommandList {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ResourceGroupName,

        [Parameter(Mandatory = $true)]
        [array]$Commands
    )

    foreach ($cmd in $Commands) {
        Invoke-CacheCommand `
            -ResourceGroupName $ResourceGroupName `
            -Command $cmd
    }
}

function View-CommandResult {
    param(
        [Parameter(Mandatory = $true)]
        [pscustomobject]$Result
    )

    if ($Result.status -eq "error") {
        Write-Host $Result.message -ForegroundColor Red
    }
    else {
        Write-Host $Result.message
    }
}

function ExecuteCommands{
    $EnvironmentName =  Get-EnvName -EnvironmentName $EnvironmentName
    #Get Vms with Given Caches
    $vms = Get-VmsWithCache -EnvironmentName $EnvironmentName -Name $Name

    #Check If VM Exists or not
    if (-not $vms -or $vms.Count -eq 0) {
        throw "No Server Exists With Given Cache"
    }

    #Get Private IPs of Vms (VM Name, CachesList, NicName, PrivateIps, OsType)
    $privateIps = Get-VMPrivateIps -VMs $vms -ResourceGroupName $EnvironmentName  -Servers $Servers

    #Get Cache Groups (CacheName, FirstVM, OsType, PrivateIps)
    $cacheGroups = Get-CacheGrouping -VMInfo $privateIps

    #Get Commands to be Executed on Azure
    $commands = Get-Commands -cacheGroups $cacheGroups
    
    #Execute Commands
    $results = Invoke-CacheCommandList -ResourceGroupName $EnvironmentName -Commands $commands

    $results | ForEach-Object { View-CommandResult -Result $_ }
}

try {
    if (-not (Get-AzContext)) {
        Connect-AzAccount
        if (Get-AzContext) {
            Write-Output "Stopping Cache..."
            ExecuteCommands
        }
    }
    else {
        Write-Output "Stopping Cache..."
        ExecuteCommands
    }
}
catch {
    Write-Error $($_.Exception.Message)

    Write-Error "Coudn't Stop Cache'"
}