<#
.SYNOPSIS
    Windows DNS Server Prometheus Metrics Exporter
.DESCRIPTION
    Prometheus exporter for Windows DNS Server - query statistics, response
    codes, zone counts, dynamic updates, recursion stats, cache performance,
    zone transfer tracking, and DNSSEC metrics via PowerShell. Exports metrics
    as Prometheus-compatible text format.
.PARAMETER Mode
    Output mode: 'stdout' (default), 'textfile', or 'http'
.PARAMETER Port
    HTTP port for http mode (default: 9536)
.PARAMETER TextfileDir
    Directory for textfile collector output (default: C:\ProgramData\node_exporter)
.PARAMETER InstallScheduledTask
    Switch to create a scheduled task for automatic execution
.PARAMETER TaskIntervalMinutes
    Interval in minutes for the scheduled task (default: 2)
.NOTES
    Author:  Phil Connor
    Contact: contact@mylinux.work
    Website: https://mylinux.work
    License: MIT
    Version: 1.0

    Metrics Exported:
      Core Status:
        - windows_dns_up
        - windows_dns_exporter_info{version}

      Query Statistics:
        - windows_dns_queries_total
        - windows_dns_responses_total
        - windows_dns_queries_per_second
        - windows_dns_recursive_queries_total
        - windows_dns_recursive_query_failures_total
        - windows_dns_tcp_queries_total
        - windows_dns_udp_queries_total

      Response Codes:
        - windows_dns_response_success_total
        - windows_dns_response_servfail_total
        - windows_dns_response_nxdomain_total
        - windows_dns_response_nxrrset_total
        - windows_dns_response_refused_total
        - windows_dns_response_formerr_total

      Dynamic Updates:
        - windows_dns_dynamic_updates_total
        - windows_dns_dynamic_updates_rejected_total
        - windows_dns_secure_updates_total

      Zones:
        - windows_dns_zones_total
        - windows_dns_zone_transfers_received_total
        - windows_dns_zone_transfers_sent_total
        - windows_dns_zone_transfer_failures_total

      Cache:
        - windows_dns_cache_hits_total
        - windows_dns_cache_flushes_total

      DNSSEC:
        - windows_dns_dnssec_signed_zones_total

      Exporter:
        - windows_dns_exporter_duration_seconds
        - windows_dns_exporter_last_run_timestamp
#>

param(
    [ValidateSet('stdout', 'textfile', 'http')]
    [string]$Mode = 'stdout',

    [int]$Port = 9536,

    [string]$TextfileDir = 'C:\ProgramData\node_exporter',

    [switch]$InstallScheduledTask,

    [int]$TaskIntervalMinutes = 2
)

if ($InstallScheduledTask) {
    $taskName = "WindowsDnsExporter"
    $existingTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue

    if (-not $existingTask) {
        $taskAction = New-ScheduledTaskAction -Execute "powershell.exe" -Argument "-NoProfile -ExecutionPolicy Bypass -File `"$($MyInvocation.MyCommand.Path)`" -Mode textfile"

        if (-not $TaskIntervalMinutes -or $TaskIntervalMinutes -le 0) {
            throw "TaskIntervalMinutes must be a positive integer"
        }

        $taskTrigger = New-ScheduledTaskTrigger -Once -At (Get-Date).AddMinutes(1) -RepetitionInterval (New-TimeSpan -Minutes $TaskIntervalMinutes) -RepetitionDuration (New-TimeSpan -Days 365)
        $taskPrincipal = New-ScheduledTaskPrincipal -UserId "SYSTEM" -LogonType ServiceAccount -RunLevel Highest

        try {
            Write-Host "Creating scheduled task: $taskName"
            Register-ScheduledTask -TaskName $taskName -Action $taskAction -Trigger $taskTrigger -Principal $taskPrincipal -Description "Exports Windows DNS Server metrics for Prometheus every $TaskIntervalMinutes minutes"

            $createdTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue
            if (-not $createdTask) {
                throw "Failed to verify scheduled task creation"
            }
            Write-Host "Successfully created scheduled task: $taskName" -ForegroundColor Green
        } catch {
            Write-Error "Failed to create scheduled task: $($_.Exception.Message)"
            throw
        }
    } else {
        Write-Host "Scheduled task '$taskName' already exists, skipping creation"
    }
}

$ErrorActionPreference = 'SilentlyContinue'

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

function Get-UnixTimestamp {
    [int][double]::Parse((Get-Date -UFormat '%s'))
}

function Format-MetricValue {
    param([double]$Value, [int]$Decimals = 2)
    [math]::Round($Value, $Decimals)
}

# ============================================================================
# DNS QUERY STATISTICS
# ============================================================================

function Get-DnsQueryMetrics {
    $sb = [System.Text.StringBuilder]::new()

    try {
        $stats = Get-DnsServerStatistics -ErrorAction Stop
        $queryStats = $stats.QueryStatistics
        $recursionStats = $stats.RecursionStatistics

        # --- Total queries ---
        [void]$sb.AppendLine('# HELP windows_dns_queries_total Total DNS queries received')
        [void]$sb.AppendLine('# TYPE windows_dns_queries_total counter')
        [void]$sb.AppendLine("windows_dns_queries_total $($queryStats.TotalQueries)")
        [void]$sb.AppendLine('')

        # --- Total responses ---
        [void]$sb.AppendLine('# HELP windows_dns_responses_total Total DNS responses sent')
        [void]$sb.AppendLine('# TYPE windows_dns_responses_total counter')
        [void]$sb.AppendLine("windows_dns_responses_total $($queryStats.TotalResponses)")
        [void]$sb.AppendLine('')

        # --- Queries per second (instantaneous) ---
        [void]$sb.AppendLine('# HELP windows_dns_queries_per_second Current query rate')
        [void]$sb.AppendLine('# TYPE windows_dns_queries_per_second gauge')
        $qps = 0
        try {
            $counter = Get-Counter '\DNS\Total Query Received/sec' -ErrorAction Stop
            $qps = Format-MetricValue $counter.CounterSamples[0].CookedValue
        } catch {}
        [void]$sb.AppendLine("windows_dns_queries_per_second $qps")
        [void]$sb.AppendLine('')

        # --- UDP vs TCP ---
        [void]$sb.AppendLine('# HELP windows_dns_udp_queries_total Total UDP queries received')
        [void]$sb.AppendLine('# TYPE windows_dns_udp_queries_total counter')
        [void]$sb.AppendLine("windows_dns_udp_queries_total $($queryStats.UdpQueries)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_tcp_queries_total Total TCP queries received')
        [void]$sb.AppendLine('# TYPE windows_dns_tcp_queries_total counter')
        [void]$sb.AppendLine("windows_dns_tcp_queries_total $($queryStats.TcpQueries)")
        [void]$sb.AppendLine('')

        # --- Recursive queries ---
        [void]$sb.AppendLine('# HELP windows_dns_recursive_queries_total Total recursive queries')
        [void]$sb.AppendLine('# TYPE windows_dns_recursive_queries_total counter')
        [void]$sb.AppendLine("windows_dns_recursive_queries_total $($recursionStats.TotalQueriesRecursed)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_recursive_query_failures_total Total recursive query failures')
        [void]$sb.AppendLine('# TYPE windows_dns_recursive_query_failures_total counter')
        [void]$sb.AppendLine("windows_dns_recursive_query_failures_total $($recursionStats.TotalRecursionFailures)")
        [void]$sb.AppendLine('')

    } catch {
        Write-Warning "Failed to collect DNS query statistics: $_"
    }

    $sb.ToString()
}

# ============================================================================
# RESPONSE CODE METRICS
# ============================================================================

function Get-DnsResponseCodeMetrics {
    $sb = [System.Text.StringBuilder]::new()

    try {
        $stats = Get-DnsServerStatistics -ErrorAction Stop
        $queryStats = $stats.QueryStatistics

        [void]$sb.AppendLine('# HELP windows_dns_response_success_total Total successful responses (NOERROR)')
        [void]$sb.AppendLine('# TYPE windows_dns_response_success_total counter')
        [void]$sb.AppendLine("windows_dns_response_success_total $($queryStats.TotalResponses - $queryStats.ServerFailureResponses - $queryStats.NameErrorResponses - $queryStats.RefusedResponses - $queryStats.FormErrorResponses)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_response_servfail_total Total SERVFAIL responses')
        [void]$sb.AppendLine('# TYPE windows_dns_response_servfail_total counter')
        [void]$sb.AppendLine("windows_dns_response_servfail_total $($queryStats.ServerFailureResponses)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_response_nxdomain_total Total NXDOMAIN responses')
        [void]$sb.AppendLine('# TYPE windows_dns_response_nxdomain_total counter')
        [void]$sb.AppendLine("windows_dns_response_nxdomain_total $($queryStats.NameErrorResponses)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_response_nxrrset_total Total NXRRSET responses')
        [void]$sb.AppendLine('# TYPE windows_dns_response_nxrrset_total counter')
        $nxrrset = 0
        try { $nxrrset = $queryStats.NxRrsetResponses } catch {}
        [void]$sb.AppendLine("windows_dns_response_nxrrset_total $nxrrset")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_response_refused_total Total REFUSED responses')
        [void]$sb.AppendLine('# TYPE windows_dns_response_refused_total counter')
        [void]$sb.AppendLine("windows_dns_response_refused_total $($queryStats.RefusedResponses)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_response_formerr_total Total FORMERR responses')
        [void]$sb.AppendLine('# TYPE windows_dns_response_formerr_total counter')
        [void]$sb.AppendLine("windows_dns_response_formerr_total $($queryStats.FormErrorResponses)")
        [void]$sb.AppendLine('')

    } catch {
        Write-Warning "Failed to collect DNS response code metrics: $_"
    }

    $sb.ToString()
}

# ============================================================================
# DYNAMIC UPDATE METRICS
# ============================================================================

function Get-DnsDynamicUpdateMetrics {
    $sb = [System.Text.StringBuilder]::new()

    try {
        $stats = Get-DnsServerStatistics -ErrorAction Stop
        $updateStats = $stats.UpdateStatistics

        [void]$sb.AppendLine('# HELP windows_dns_dynamic_updates_total Total dynamic updates received')
        [void]$sb.AppendLine('# TYPE windows_dns_dynamic_updates_total counter')
        [void]$sb.AppendLine("windows_dns_dynamic_updates_total $($updateStats.DynamicUpdatesReceived)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_dynamic_updates_rejected_total Total dynamic updates rejected')
        [void]$sb.AppendLine('# TYPE windows_dns_dynamic_updates_rejected_total counter')
        [void]$sb.AppendLine("windows_dns_dynamic_updates_rejected_total $($updateStats.DynamicUpdatesRejected)")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_secure_updates_total Total secure dynamic updates received')
        [void]$sb.AppendLine('# TYPE windows_dns_secure_updates_total counter')
        [void]$sb.AppendLine("windows_dns_secure_updates_total $($updateStats.SecureUpdatesReceived)")
        [void]$sb.AppendLine('')

    } catch {
        Write-Warning "Failed to collect DNS dynamic update metrics: $_"
    }

    $sb.ToString()
}

# ============================================================================
# ZONE METRICS
# ============================================================================

function Get-DnsZoneMetrics {
    $sb = [System.Text.StringBuilder]::new()

    try {
        $zones = Get-DnsServerZone -ErrorAction Stop | Where-Object { $_.ZoneType -ne 'Forwarder' -and $_.ZoneName -ne 'TrustAnchors' }

        [void]$sb.AppendLine('# HELP windows_dns_zones_total Total number of DNS zones')
        [void]$sb.AppendLine('# TYPE windows_dns_zones_total gauge')
        [void]$sb.AppendLine("windows_dns_zones_total $($zones.Count)")
        [void]$sb.AppendLine('')

        # Zone transfer stats from server statistics
        $stats = Get-DnsServerStatistics -ErrorAction Stop
        $zoneTransferStats = $stats.ZoneTransferStatistics

        [void]$sb.AppendLine('# HELP windows_dns_zone_transfers_received_total Total zone transfers received (AXFR/IXFR in)')
        [void]$sb.AppendLine('# TYPE windows_dns_zone_transfers_received_total counter')
        $xfrIn = 0
        try { $xfrIn = $zoneTransferStats.TransferRequestsReceived } catch {}
        [void]$sb.AppendLine("windows_dns_zone_transfers_received_total $xfrIn")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_zone_transfers_sent_total Total zone transfers sent (AXFR/IXFR out)')
        [void]$sb.AppendLine('# TYPE windows_dns_zone_transfers_sent_total counter')
        $xfrOut = 0
        try { $xfrOut = $zoneTransferStats.SuccessfulTransfersSent } catch {}
        [void]$sb.AppendLine("windows_dns_zone_transfers_sent_total $xfrOut")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_zone_transfer_failures_total Total zone transfer failures')
        [void]$sb.AppendLine('# TYPE windows_dns_zone_transfer_failures_total counter')
        $xfrFail = 0
        try { $xfrFail = $zoneTransferStats.FailedTransfers } catch {}
        [void]$sb.AppendLine("windows_dns_zone_transfer_failures_total $xfrFail")
        [void]$sb.AppendLine('')

    } catch {
        Write-Warning "Failed to collect DNS zone metrics: $_"
    }

    $sb.ToString()
}

# ============================================================================
# CACHE METRICS
# ============================================================================

function Get-DnsCacheMetrics {
    $sb = [System.Text.StringBuilder]::new()

    try {
        $stats = Get-DnsServerStatistics -ErrorAction Stop
        $cacheStats = $stats.CacheStatistics

        [void]$sb.AppendLine('# HELP windows_dns_cache_hits_total Total cache hits')
        [void]$sb.AppendLine('# TYPE windows_dns_cache_hits_total counter')
        $cacheHits = 0
        try { $cacheHits = $cacheStats.CacheHits } catch {}
        [void]$sb.AppendLine("windows_dns_cache_hits_total $cacheHits")
        [void]$sb.AppendLine('')

        [void]$sb.AppendLine('# HELP windows_dns_cache_flushes_total Total cache flushes')
        [void]$sb.AppendLine('# TYPE windows_dns_cache_flushes_total counter')
        $cacheFlushes = 0
        try { $cacheFlushes = $cacheStats.CacheFlushes } catch {}
        [void]$sb.AppendLine("windows_dns_cache_flushes_total $cacheFlushes")
        [void]$sb.AppendLine('')

    } catch {
        Write-Warning "Failed to collect DNS cache metrics: $_"
    }

    $sb.ToString()
}

# ============================================================================
# DNSSEC METRICS
# ============================================================================

function Get-DnsDnssecMetrics {
    $sb = [System.Text.StringBuilder]::new()

    try {
        $zones = Get-DnsServerZone -ErrorAction Stop | Where-Object { $_.ZoneType -ne 'Forwarder' -and $_.ZoneName -ne 'TrustAnchors' }
        $signedCount = 0
        foreach ($zone in $zones) {
            try {
                $signing = Get-DnsServerDnsSecZoneSetting -ZoneName $zone.ZoneName -ErrorAction Stop
                if ($signing.IsSigned) { $signedCount++ }
            } catch {}
        }

        [void]$sb.AppendLine('# HELP windows_dns_dnssec_signed_zones_total Number of DNSSEC-signed zones')
        [void]$sb.AppendLine('# TYPE windows_dns_dnssec_signed_zones_total gauge')
        [void]$sb.AppendLine("windows_dns_dnssec_signed_zones_total $signedCount")
        [void]$sb.AppendLine('')

    } catch {
        Write-Warning "Failed to collect DNSSEC metrics: $_"
    }

    $sb.ToString()
}

# ============================================================================
# COLLECT ALL METRICS
# ============================================================================

function Get-AllMetrics {
    $scriptStart = Get-Date
    $sb = [System.Text.StringBuilder]::new()

    # Exporter up
    [void]$sb.AppendLine('# HELP windows_dns_up Exporter status (1=up, 0=down)')
    [void]$sb.AppendLine('# TYPE windows_dns_up gauge')
    [void]$sb.AppendLine('windows_dns_up 1')
    [void]$sb.AppendLine('')

    # Exporter info
    [void]$sb.AppendLine('# HELP windows_dns_exporter_info Exporter version information')
    [void]$sb.AppendLine('# TYPE windows_dns_exporter_info gauge')
    [void]$sb.AppendLine('windows_dns_exporter_info{version="1.0"} 1')
    [void]$sb.AppendLine('')

    # Collect all metric sections
    [void]$sb.Append((Get-DnsQueryMetrics))
    [void]$sb.Append((Get-DnsResponseCodeMetrics))
    [void]$sb.Append((Get-DnsDynamicUpdateMetrics))
    [void]$sb.Append((Get-DnsZoneMetrics))
    [void]$sb.Append((Get-DnsCacheMetrics))
    [void]$sb.Append((Get-DnsDnssecMetrics))

    # Exporter runtime
    $scriptEnd = Get-Date
    $duration = Format-MetricValue ($scriptEnd - $scriptStart).TotalSeconds
    $timestamp = Get-UnixTimestamp

    [void]$sb.AppendLine('# HELP windows_dns_exporter_duration_seconds Time to generate all metrics')
    [void]$sb.AppendLine('# TYPE windows_dns_exporter_duration_seconds gauge')
    [void]$sb.AppendLine("windows_dns_exporter_duration_seconds $duration")
    [void]$sb.AppendLine('')
    [void]$sb.AppendLine('# HELP windows_dns_exporter_last_run_timestamp Unix timestamp of last successful run')
    [void]$sb.AppendLine('# TYPE windows_dns_exporter_last_run_timestamp gauge')
    [void]$sb.AppendLine("windows_dns_exporter_last_run_timestamp $timestamp")
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# HTTP SERVER MODE
# ============================================================================

function Start-HttpServer {
    param([int]$ListenPort)

    $prefix = "http://+:$ListenPort/"
    $listener = [System.Net.HttpListener]::new()
    $listener.Prefixes.Add($prefix)

    try {
        $listener.Start()
        Write-Host "Starting Windows DNS exporter on port $ListenPort..." -ForegroundColor Green
        Write-Host "Metrics available at http://localhost:$ListenPort/metrics"

        while ($listener.IsListening) {
            $context  = $listener.GetContext()
            $request  = $context.Request
            $response = $context.Response

            if ($request.Url.AbsolutePath -eq '/metrics') {
                $metrics = Get-AllMetrics
                $buffer  = [System.Text.Encoding]::UTF8.GetBytes($metrics)
                $response.ContentType = 'text/plain; version=0.0.4; charset=utf-8'
            }
            else {
                $html = @"
<!DOCTYPE html>
<html>
<head><title>Windows DNS Exporter v1.0</title></head>
<body>
<h1>Windows DNS Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<h2>Collectors</h2>
<ul>
<li>Query statistics (total, UDP, TCP, recursive)</li>
<li>Response codes (SUCCESS, SERVFAIL, NXDOMAIN, REFUSED, FORMERR)</li>
<li>Dynamic updates (received, rejected, secure)</li>
<li>Zone statistics (count, transfers)</li>
<li>Cache (hits, flushes)</li>
<li>DNSSEC (signed zone count)</li>
</ul>
</body>
</html>
"@
                $buffer = [System.Text.Encoding]::UTF8.GetBytes($html)
                $response.ContentType = 'text/html; charset=utf-8'
            }

            $response.ContentLength64 = $buffer.Length
            $response.OutputStream.Write($buffer, 0, $buffer.Length)
            $response.OutputStream.Close()
        }
    }
    catch {
        Write-Error "HTTP server error: $_"
        Write-Error "If access denied, run: netsh http add urlacl url=http://+:$ListenPort/ user=Everyone"
    }
    finally {
        if ($listener.IsListening) {
            $listener.Stop()
        }
    }
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

switch ($Mode) {
    'http' {
        Start-HttpServer -ListenPort $Port
    }
    'textfile' {
        $OutputFile = Join-Path $TextfileDir 'windows_dns.prom'

        $outputDir = Split-Path $OutputFile -Parent
        if (-not (Test-Path $outputDir)) {
            New-Item -Path $outputDir -ItemType Directory -Force | Out-Null
        }

        $tempFile = Join-Path $outputDir ".windows_dns_metrics.$PID.tmp"

        try {
            $metrics = Get-AllMetrics
            $metrics | Out-File -FilePath $tempFile -Encoding utf8 -NoNewline

            $lineCount = ($metrics -split "`n").Count
            if ($lineCount -lt 10) {
                Remove-Item -Path $tempFile -Force -ErrorAction SilentlyContinue
                Write-Error "Metrics file too small ($lineCount lines), keeping previous"
                exit 1
            }

            Move-Item -Path $tempFile -Destination $OutputFile -Force
            Write-Host "Metrics written to $OutputFile ($lineCount lines)" -ForegroundColor Green
        }
        catch {
            Remove-Item -Path $tempFile -Force -ErrorAction SilentlyContinue
            Write-Error "Failed to generate metrics: $_"
            exit 1
        }
    }
    default {
        Get-AllMetrics | Write-Output
    }
}
