# Variables
$regPath = "HKCU:\Software\Zscaler\App"

# Lists/Arrays
$requiredServices = @("ZSAService", "ZSATunnel")
$optionalServices = @("ZSATrayManager", "ZSAUpm")

# Function to check if Zscaler is installed
function Test-ZscalerInstalled {

    $zscalerFiles = @(
        "C:\Program Files\Zscaler\ZSAService\ZSAService.exe",
        "C:\Program Files\Zscaler\ZSATunnel\ZSATunnel.exe",
        "C:\Program Files\Zscaler\ZSATrayManager\ZSATrayManager.exe",
        "C:\Program Files\Zscaler\ZSAUpm\ZSAUpm.exe"
    )

    foreach ($file in $zscalerFiles) {

        Write-Verbose "Checking Zscaler installation path: $file"

        if (Test-Path $file) {
            Write-Verbose "Zscaler installation found."
            return $true
        }
    }

    Write-Warning "Zscaler installation not found."
    return $false
}

# Function to check if a service exists
function Test-ServiceExists {
    param (
        [string]$serviceName
    )

    $service = Get-Service -Name $serviceName -ErrorAction SilentlyContinue

    if ($service) {
        Write-Verbose "Service exists: $serviceName"
        return $true
    }

    Write-Warning "Service does not exist: $serviceName"
    return $false
}

# Function to get service status
function Get-ServiceStatus {
    param (
        [string]$serviceName
    )

    try {

        $service = Get-Service -Name $serviceName -ErrorAction Stop

        if ($service.Status -eq 'Running') {
            Write-Verbose "$serviceName is Running"
            return $true
        }

        Write-Warning "$serviceName is $($service.Status)"
        return $false
    }
    catch {
        Write-Warning "$serviceName not found"
        return $false
    }
}

# Function to retrieve registry values
function Get-RegistryValue {
    param (
        [string]$path,
        [string]$name
    )

    try {
        return (Get-ItemProperty -Path $path -Name $name -ErrorAction Stop).$name.Trim()
    }
    catch {
        Write-Warning "$name not found in $path"
        return "Not Found"
    }
}

# Function to verify Zscaler authentication
function Test-ZscalerAuthentication {

    $url = "https://ip.zscaler.com/"
    $attempts = 3

    for ($i = 1; $i -le $attempts; $i++) {

        try {

            $response = Invoke-WebRequest `
                -Uri $url `
                -UseBasicParsing `
                -TimeoutSec 30

            if (
                $response.StatusCode -eq 200 -and
                (
                    $response.Content -match "Zscaler" -or
                    $response.Content -match "Your request is arriving"
                )
            ) {
                return $true
            }

        }
        catch {
            Write-Warning "Authentication test attempt $i failed."
        }

        Start-Sleep -Seconds 3
    }

    return $false
}

# Main Function
function Test-ZscalerStatus {

    $zscalerInstalled = Test-ZscalerInstalled

    if (-not $zscalerInstalled) {
        Write-Warning "Zscaler installation not detected."
        return $false
    }

    $allServicesRunning = $true

    foreach ($serviceName in $requiredServices) {

        $status = Get-ServiceStatus -serviceName $serviceName

        if (-not $status) {
            $allServicesRunning = $false
        }
    }

    $znwState = Get-RegistryValue -path $regPath -name "ZNW_State"
    $zpaState = Get-RegistryValue -path $regPath -name "ZPA_State"
    $zwsState = Get-RegistryValue -path $regPath -name "ZWS_State"
    $zpaAuthState = Get-RegistryValue -path $regPath -name "ZPAAuth_State"

    $zpaStateValid = $zpaState -eq "TUNNEL_FORWARDING"
    $zwsStateValid = $zwsState -eq "TUNNEL_FORWARDING"
    $zpaAuthValid = $zpaAuthState -eq "AUTHENTICATED"

    Write-Host ""
    Write-Host "===== ZSCALER STATUS ====="
    Write-Host "Installed      : $zscalerInstalled"
    Write-Host "ZNW State      : $znwState"
    Write-Host "ZPA State      : $zpaState"
    Write-Host "ZWS State      : $zwsState"
    Write-Host "ZPA Auth State : $zpaAuthState"
    Write-Host "Services OK    : $allServicesRunning"
    Write-Host "=========================="
    Write-Host ""

    $isZscalerOperational = (
        $zscalerInstalled -and
        $allServicesRunning -and
        $zpaAuthValid -and
        $zpaStateValid -and
        $zwsStateValid
    )

    return $isZscalerOperational
}

# Run Test
Test-ZscalerStatus