Dev.Core

1.0.1

dev.core.dns.psm1

Import-Module "$PSScriptRoot\dev.core.utils.psm1" -DisableNameChecking;

$hostsFile = "${env:SystemRoot}\system32\drivers\etc\hosts";

[regex]$hostEntryRegex = "(?<ip>[\d\.]+)\s+(?<host>[\w\d\.\-_]+)";

#
# Private functions
#
function Update-HostsFile
{
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)][object[]]$Value,
        [Parameter(Mandatory=$true, ParameterSetName="Update")][switch]$Update, 
        [Parameter(Mandatory=$true, ParameterSetName="Append")][switch]$Append);

    $backupHostsFile = "${env:SystemRoot}\system32\drivers\etc\hosts.$(New-Guid).bak";
            
    # Backup the hosts file before making any change
    Write-Verbose "Copying hosts file to ${backupHostsFile}...";
    Copy-Item -Path $script:hostsFile -Destination $backupHostsFile;

    $error = "The ${script:hostsFile} file is not updated successfully, please restore hosts file using back file ${backupHostsFile}.";

    try
    {
        if ($PSCmdlet.ParameterSetName -eq "Update")
        {
            $expected = $Value;
            Write-Verbose "Updating hosts file to ${script:hostsFile}...";
            Set-Content -Path $script:hostsFile -Value $Value;            
        }
        else # Append
        {
            $expected = (Get-Content -Path $script:hostsFile) + $Value;
            Write-Verbose "Appending hosts file to ${script:hostsFile}...";
            Add-Content -Path $script:hostsFile -Value $appendLines;
        }

        Write-Verbose "Waiting 1 second for the change to be picked up by windows";
        Start-Sleep -Seconds 1;
    }
    catch
    {
        Write-Error $_;
        throw $error;
    }

    # Read content from hosts again to make sure the change is successful.
    # Here do not try to restore the hosts file as most likely the restore will fail again.
    $newValue =  Get-Content -Path $script:hostsFile;
    if (($expected -join "`r`n") -ne ($newValue -join "`r`n"))
    {
        Write-Error "The new hosts file doesn't match expected content";
        throw $error;
    }
    else
    {
        # If the new content matches expected content, then delete the backup file.
        Write-Verbose "Deleting backup file ${backupHostsFile}...";
        Remove-Item -Path $backupHostsFile;
    }
}


#
# Public functions
#

function New-LocalDnsEntry
{
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)][string]$HostName, 
        [Parameter(Mandatory=$true)][string]$IPAddress);

    $entry = [PSCustomObject]@{ 
        HostName = $HostName; 
        IPAddress = $IPAddress 
    };
    
    Write-Output $entry;
}

function Get-LocalDnsEntry
{
    [CmdletBinding()]
    param([string]$HostName);

    $lines = Get-Content -Path $script:hostsFile;
    
    for ($i = 0; $i -lt $lines.Length; $i++)
    {
        $line = $lines[$i];

        if ($line.Trim().StartsWith("#"))
        {
            continue;
        }

        $match = $hostEntryRegex.Match($line);
        if ($match.Success)
        {
            $capturedHost = $match.Groups["host"].value;
            $capturedIp = $match.Groups["ip"].Value;

            if (!$HostName -or $capturedHost -eq $HostName)
            {
                New-LocalDnsEntry -HostName $capturedHost -IPAddress $capturedIp;
            }
        }
    }
}

function Set-LocalDnsEntry
{
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true, ValueFromPipelineByPropertyName=$true)][string]$HostName, 
        [Parameter(Mandatory=$true, ValueFromPipelineByPropertyName=$true)][string]$IPAddress);

    begin
    {
        if (!(Test-RunAsAdmin))
        {
            throw "The Set-LocalDnsEntry command must be executed with administrator privilege."
        }

        $lines = Get-Content -Path $script:hostsFile;

        $updateLines = $lines;
        $appendLines = @();

        $shouldUpdate = $false;
        $shouldAppend = $false;
    }

    process
    {
        $hostEntry = "${IPAddress}`t${HostName}";
        $found = $false;

        for ($i = 0; $i -lt $updateLines.Length; $i++)
        {
            $line = $updateLines[$i];

            if ($line.Trim().StartsWith("#"))
            {
                continue;
            }

            $match = $hostEntryRegex.Match($line);
            if ($match.Success)
            {
                $capturedHost = $match.Groups["host"].value;
                $capturedIP = $match.Groups["ip"].value;
                if ($capturedHost -eq $HostName)
                {
                    if ($capturedIP -ne $IPAddress)
                    {
                        Write-Verbose "Updating DNS entry from `"$line`" to `"$hostEntry`"...";
                       
                        $updateLines[$i] = $hostEntry;
                        $shouldUpdate = $true;
                    }
                    
                    $found = $true;
                }
            }
        }
        
        if (!$found)
        {
            Write-Verbose "Appending DNS entry `"$hostEntry`"...";
            $appendLines += $hostEntry;
            $shouldAppend = $true;
        }
    }

    end
    {
        if ($shouldUpdate)
        {
            Update-HostsFile -Value ($updateLines + $appendLines) -Update;
        }
        elseif ($shouldAppend)
        {
            Update-HostsFile -Value $appendLines -Append;
        }
        else
        {
            Write-Warning "Skip updating ${script:hostsFile} file as there is no change detected";
        }
    }
}

function Remove-LocalDnsEntry
{
    [CmdletBinding(SupportsShouldProcess=$true)]
    param([Parameter(Mandatory=$true, ValueFromPipelineByPropertyName=$true)][string]$HostName);

    begin
    {
    Confirm-RunAsAdmin -Command $MyInvocation.MyCommand;

        $lines = Get-Content -Path $script:hostsFile;
        $shouldUpdate = $false;     
    }

    process
    {
        for ($i = 0; $i -lt $lines.Length; $i++)
        {
            $line = $lines[$i];

            if ($line -ne $null -and !$line.Trim().StartsWith("#"))
            {
                $match = $hostEntryRegex.Match($line);
                if ($match.Success)
                {
                    $capturedHost = $match.Groups["host"].value;
                    if ($capturedHost -eq $HostName)
                    {
                        Write-Verbose "Removing DNS entry from `"$line`"...";
                        $lines[$i] = $null; # Set the line to $null for deletion.
                        $shouldUpdate = $true;
                    }
                }
            }
        }
    }

    end
    {
        if ($shouldUpdate)
        {
            Update-HostsFile -Value ($lines | ? { $_ -ne $null }) -Update;
        }
        else
        {
            Write-Warning "Skip updating ${script:hostsFile} file as there is no change detected";
        }
    }
}

#
# Module Initialization
#

Export-ModuleMember -Function New-LocalDnsEntry;
Export-ModuleMember -Function Get-LocalDnsEntry;
Export-ModuleMember -Function Set-LocalDnsEntry;
Export-ModuleMember -Function Remove-LocalDnsEntry;