#
# ConnectivityCheck.ps1
#
param (
	[parameter(Mandatory=$true)]
	[string]
	$NameResolutionFileName,

    [parameter(Mandatory=$true)]
    [string]
    $TcpEndpointsFileName,
    
    [parameter(Mandatory=$false)]
    [string]
    $OutDir = $Env:Temp,
    
    [parameter(Mandatory=$false)]
    [string]
    $CheckHostsFile
 )

 $BlnCheckHostsFile = $false
 if ($CheckHostsFile) {
    if ($CheckHostsFile -ine "false" -and $CheckHostFile -ine "no" -and $CheckHostFile -ne "0") {
       $BlnCheckHostsFile = $true
    }
 }

 # Functions
 #

# Override the built-in cmdlet with a custom version
function Write-ErrorMsg($Message) {
    [Console]::ForegroundColor = 'red'
    Write-Output $Message
    [Console]::ResetColor()
}

function Convert-ByteArrayToHex {
    [cmdletbinding()]
    param(
        [parameter(Mandatory=$true)]
        [Byte[]]
        $Bytes
    )
    #$HexString = [System.Text.StringBuilder]::new($Bytes.Length * 2)
    $HexString = New-Object System.Text.StringBuilder($Bytes.Length * 2)
    ForEach($Byte in $Bytes){
        $HexString.AppendFormat("{0:x2}", $Byte) | Out-Null
    }
    $HexString.ToString()
}

function Convert-HexToByteArray {
	[cmdletbinding()]
	param(
		[parameter(Mandatory=$true)]
		[String]
		$HexString)
#	$Bytes = [byte[]]::new($HexString.Length / 2)
	$Bytes = New-Object byte[]($HexString.Length / 2)
	For($i=0; $i -lt $HexString.Length; $i+=2){
		$Bytes[$i/2] = [convert]::ToByte($HexString.Substring($i, 2), 16)
	}
	$Bytes
}

function IsPortOpen {
	param(
		[parameter(Mandatory=$true)]
		[string]
		$HostName,
		[parameter(Mandatory=$true)]
		[int]
		$Port,
		[parameter(Mandatory=$false)]
		[int]
		$Timeout = 5000)
	
	$Result = $null

	# REMOVED: TCPClient's ConnectAsync not available in older powershell
	# $Client = New-Object System.Net.Sockets.TCPClient
	# $Timeout = 
    # $IsOpen = $Client.ConnectAsync($HostName, $Port).Wait($Timeout)

	$Client = New-Object System.Net.Sockets.TCPClient
	$IsOpen = $false
	try {
		$AsyncResult = $Client.BeginConnect($HostName, $Port, $null, $null)
		if (!$AsyncResult.AsyncWaitHandle.WaitOne($Timeout, $false))
		{
			throw [System.TimeoutException]::new();
		}
		$Client.EndConnect($AsyncResult) | Out-Null
		$IsOpen = $Client.Connected
	}
	catch {
		$IsOpen = $false
	}
	finally {
		$Client.Close();
	}

	if ($IsOpen) {
		$Result = "Success"
	}
	else {
		$Result = "Fail"
	}
	"$($Result): $($HostName):$($Port)"
}

function Read-Hosts-File-To-Hash {
	Param (
		[string] $FileName,
		[int] $IncludeCommentedLines)

	$HostHash = @{}
	if (-not (Test-Path -Path $FileName)) {
		Write-ErrorMsg "Error: file not found: $($FileName)"
		return $null
	}
	# $Lines = Get-Content -Path $FileName
	$Lines = [IO.File]::ReadAllLines($FileName)
	foreach ($Line in $Lines) {
		$Line = $Line.Trim()
		if ($Line.Length -lt 9) {
			continue
		}
		$Chars = $Line.ToCharArray()
		if ($Chars[0] -eq '#') {
			if ($IncludeCommentedLines -eq 0) {
				continue
			}
		}
		$Line = $Line.Trim('#')
		if ($Line.Length -lt 9) {
			continue
		}
		$ix = $Line.IndexOf('#')
		if ($ix -gt 0) {
			if ($ix -lt 9) {
				continue
			}
			$Line = $Line.Substring(0, $ix)
		}
		$Words = ($Line -split ' ')
		if ($Words.Count -lt 2) {
			continue
		}
		$Ip = $Words[0]
		try {
			[IPAddress]$Ip > $null
		}
		catch {
			continue
		}
		$Names = New-Object System.Collections.Generic.List[String]
		for ($i = 1; $i -lt $Words.Count; ++$i) {
			$Word = $Words[$i].Trim()
			if ($Word) {
				$Names.Add($Words[$i])
			}
		}
		if ($HostHash.ContainsKey($Ip)) {
			foreach ($Name in $Names) {
				$HostHash[$Ip].Add($Name)
			}
		}
		else {
			$HostHash.Add($Ip, $Names)
		}
	}
	$HostHash
}

function Read-TcpEndpoints-File-To-Hash {
	Param (
		[string] $FileName,
		[int] $IncludeCommentedLines)

	$TcpEndpointsHash = @{}
	if (-not (Test-Path -Path $FileName)) {
		Write-ErrorMsg "Error: file not found: $($FileName)"
		return $null
	}
	# $Lines = Get-Content -Path $FileName
	$Lines = [IO.File]::ReadAllLines($FileName)
	foreach ($Line in $Lines) {
		$Line = $Line.Trim()
		if ($Line.Length -lt 9) {
			continue
		}
		$Chars = $Line.ToCharArray()
		if ($Chars[0] -eq '#') {
			if ($IncludeCommentedLines -eq 0) {
				continue
			}
		}
		$Line = $Line.Trim('#')
		if ($Line.Length -lt 9) {
			continue
		}
		$ix = $Line.IndexOf('#')
		if ($ix -gt 0) {
			if ($ix -lt 9) {
				continue
			}
			$Line = $Line.Substring(0, $ix)
		}
		$Parts = ($Line -split ' ')
		if ($Parts.Count -lt 2) {
			continue
		}
		$Ip = $Parts[0]
		try {
			[IPAddress]$Ip > $null
		}
		catch {
			continue
		}
		$PortsCsv = $Parts[1]
		$PortsArr = ($PortsCsv -split ',')
		$Ports = New-Object System.Collections.Generic.List[String]
		for ($i = 0; $i -lt $PortsArr.Count; ++$i) {
			$Port = $PortsArr[$i]
			$Port = $Port.Trim()
			if ($Port) {
				$Ports.Add($Port)
			}
		}
		if ($TcpEndpointsHash.ContainsKey($Ip)) {
			foreach ($Port in $Ports) {
				$TcpEndpointsHash[$Ip] = $Ports
			}
		}
		else {
			$TcpEndpointsHash.Add($Ip, $Ports)
		}
	}
	$TcpEndpointsHash
}


# Set Up
#

$DateStringPrefix = Get-Date -Format "yyyy-MM-dd_HHmmss__"
if (-not (Test-Path -Path $OutDir)) {
    New-Item -ItemType Directory -Force -Path $OutDir >$null
}
$OutFileName = Join-Path -Path $OutDir -ChildPath "$($DateStringPrefix)ConnectivityCheck.log"
$ErrorActionPreference="SilentlyContinue"
Stop-Transcript | out-null
$ErrorActionPreference = "Continue"
Start-Transcript -path $OutFileName
Write-Host

# Retrieve IPs and host names from file for name resolution check.
$HostsFileName = "$($Env:windir)\system32\drivers\etc\hosts"
if (Test-Path -Path $NameResolutionFileName) {
	Write-Host "Retrieving IPs and host names from $NameResolutionFileName..."
	$NameResolutionTable = Read-Hosts-File-To-Hash $NameResolutionFileName 0
} else {
	Write-ErrorMsg "Name resolution file missing: $NameResolutionFileName"
	$NameResolutionTable = $null
}
if (Test-Path -Path $HostsFileName) {
	$NameResolutionTable = Read-Hosts-File-To-Hash $HostsFileName.ToString() 0
} else {
	Write-ErrorMsg "System hosts file missing: $HostsFileName"
	$NameResolutionTable = $null
}

# Retrieve TCP endpoints from file for connectivity check.
if (Test-Path -Path $TcpEndpointsFileName) {
	Write-Host "Retrieving TCP endpoints from $TcpEndpointsFileName..."
	$TcpEndpointMap = Read-TcpEndpoints-File-To-Hash $TcpEndpointsFileName
} else {
	Write-ErrorMsg "TCP endpoints file missing: $TcpEndpointsFileName"
	$TcpEndpointMap = $null
}


# Name Resolution Check
#

Write-Output "`r`n`r`nNAME RESOLUTION TEST ($NameResolutionFileName):`r`n"
if ($null -ne $NameResolutionTable) {
	foreach ($IpCorrect in $NameResolutionTable.Keys) {
		$Names = $NameResolutionTable[$IpCorrect]
		foreach ($Name in $Names) {
			if (!$Name) {
				continue
			}
			$AddressList = $null
			try {
				$AddressList = [System.Net.Dns]::GetHostAddresses($Name)
			}
			catch {
				Write-ErrorMsg "Error: $($Name) doesn't resolve."
				continue
			}
			$IpList = New-Object 'System.Collections.Generic.List[string]'
			foreach ($Ip in $AddressList) {
				if ($Ip.AddressFamily -ne "InterNetwork") {
					continue
				}
				if ($IpList.Contains($Ip.IPAddressToString)) {
					continue
				}
				$IpList.Add($Ip.IPAddressToString) > $null
			}
			if ($IpList.Count -eq 0) {
				Write-ErrorMsg "Error: $($Name) doesn't resolve to any IPv4 addresses."
				continue
			}
			if ($IpList.Count -gt 1) {
				Write-ErrorMsg "Error: $($Name) resolves to multiple IP addresses: $($IpList -join ', ')"
				continue
			}
			$Ip = $IpList[0]
			if ($IpCorrect -eq $Ip) {
				#Write-Output "Success: $($Name) resolves to $($IpCorrect)"
				Write-Output "Success: $($IpCorrect) resolved from $($Name)"
			}
			else {
				Write-ErrorMsg "Error: $($Name) resolves to $($Ip) instead of $($IpCorrect)."
			}
		}
	}
}


# Hosts file check
#

if ($BlnCheckHostsFile) {
    Write-Output "`r`n`r`nHOSTS FILE CHECK ($($HostsFileName)):`r`n"
    if ($null -ne $NameResolutionTable -and $null -ne $NameResolutionTable) {
        foreach ($Ip in $NameResolutionTable.Keys) {
            $SourceNames = $NameResolutionTable[$Ip]
            $SourceNamesStr = $SourceNames -join " "
            $SourceEntryStr = "$($Ip) $($SourceNamesStr)"
            if (-not $NameResolutionTable.ContainsKey($Ip)) {
                Write-ErrorMsg "Error: missing entry equivalent to '$($SourceEntryStr)'"
                continue
            }

            $MissingNames = @()
            $Names = $NameResolutionTable[$Ip]
            foreach ($SourceName in $SourceNames) {
                if (-not ($Names -contains $SourceName)) {
                    $MissingNames += $SourceName
                }
            }
            if (0 -lt $MissingNames.count) {
                Write-ErrorMsg "Error: entry or entries for $($Ip) missing the following names: $($MissingNames -join ' ')"
            } else {
                foreach ($SourceName in $SourceNames) {
                    Write-Output "Success: found $($Ip) $($SourceName)"
                }
            }
        }
    }
}


# TCP connectivity check
#

$EndPoints = New-Object System.Collections.Generic.List[Object]
Write-Output "`r`n`r`nTCP CONNECTIVITY TEST:`r`n"
if ($null -ne $TcpEndpointMap) {
	foreach ($Ip in $TcpEndpointMap.Keys)
	{
		$Ports = $TcpEndpointMap[$Ip]
		foreach ($Port in $Ports) {	
			$ThisEndPoint = New-Object PSObject -Property @{
				Host = [string]$Ip
				Port = [int]$Port
			}
			$EndPoints.Add($ThisEndPoint)
		}
	}
}
foreach ($EndPoint in $EndPoints) {
	IsPortOpen -hostName $EndPoint.Host -port $EndPoint.Port
}


# Clean up
#

Stop-Transcript