param(
    [string]$add = "", #add server to hostfile, multiple servers separated with semicolon, IP and server separated with pipeline
    [string]$remove = "", #remove server from hostfile, defined by IP address, multiple adresses separated with semicolon
    [string]$inlcudeShort = "yes", #switch to add or not to add shortname to the hostfile
    [string]$onlyShort = "no", #add only shortname to the file
    [string]$checkIP = "yes" #to check if IP is already used
)

$date = Get-Date -f yyyy-MM-dd
$originalFile = "C:\Windows\System32\drivers\etc\hosts" #location of original hosts file
$backupFile = "C:\Windows\System32\drivers\etc\hosts.original-$date" #location for copy of original hosts file, in case anything goes wrong, you can revert
$workingFile = "C:\Temp\hoststemp.txt" #temporary file to work with
try {
    Copy-Item $originalFile -Destination $backupFile -Force #copying original file to backup file, in case you need to revert
    Copy-Item $originalFile -Destination $workingFile -Force #copying original file to working file
    add-content $workingFile "`n" #adding new line to prevent adding new content right after old entry
}
catch {
    Write-host "Unable to create backup file. Cannot continue."
    exit
}

#check if provided IP is real IPv4
function isIP {
    param($ip)
    $isOK = $true
    $octets = $ip.split(".")
    if ($octets.count -ne 4) {
        Write-Host "Wrong IP, it does not have 4 octets: "$ip
        $isOK = $false
    }
    if (($octets[0] -notmatch "[0-9]") -or ($octets[1] -notmatch "[0-9]") -or ($octets[2] -notmatch "[0-9]") -or ($octets[3] -notmatch "[0-9]")) {
        Write-Host "Wrong IP, one or more octets do not contain number: "$ip
        $isOK = $false
    }
    elseif (([int]$octets[0] -gt 255) -or ([int]$octets[1] -gt 255) -or ([int]$octets[2] -gt 255) -or ([int]$octets[3] -gt 255)) {
        Write-Host "Wrong IP, one or more octets above the range: "$ip
        $isOK = $false
    }
    elseif (([int]$octets[0] -lt 1) -or ([int]$octets[1] -lt 0) -or ([int]$octets[2] -lt 0) -or ([int]$octets[3] -lt 1)) {
        Write-Host "Wrong IP, one or more octets below the range: "$ip
        $isOK = $false
    }
    return $isOK
}

#check if provided server name is FQDN
function isFQDN {
    param($name)
    $dn = $name.split(".")
    if ($dn.count -gt 1) {
        return $true
    }
    else {
        Write-host "Server name is not FQDN: "$name
        return $false
    }
}

#check if server is in hostfile
function checkServer {
    param($file, $ip)
    $content = (Get-Content $file)
    if($content | Select-String $ip -Quiet) {
        return $true
    }
    else {
        return $false
    }
}

#add server
function addServer {
    param($file, $ip, $server, $short)
    $hostn = $server.split(".")
    if ($short -eq "only") { #check if we like to include short name or add only shortname
        $hostname = ""
        $server = $hostn[0]
    }
    elseif ($short -eq "include") {
        $hostname = $hostn[0]
    }
    elseif ($short -eq "notInclude") {
        $hostname = "";
    }
    $error.clear()
    try { 
        Add-Content $file "$ip`t$server`t$hostname"
        Write-Host "Server $server has been added to hostfile"
    }
    catch {
        $_.Exception.Message
        Write-Host "Unable to add $server to hostfile"
    }
}

#remove server
function removeServer {
    param($file, $ip)
    $content = (Get-Content $file)
    try {
        $content | Where-Object {$_ -notmatch $ip} | Set-Content $file
        Write-Host "Server with IP $ip has been removed from hostfile"
    }
    catch {
        $_.Exception.Message
        Write-Host "Unable to remove server from hostfile"
    }
}

#get information of local server
function getLocalInfo {
    $record = @()
    $hostname = hostname
    $ipAddress = ((Test-Connection -ComputerName $hostname -Count 1).IPV4Address).IPAddressToString #get local IP Adress
    $DNSName = ([System.Net.Dns]::GetHostbyAddress($ipAddress)).HostName
    $record += ,$ipAddress
    $record += ,$DNSName
    return $record
}

#add servers to hostfile
if ($add -ne "") {
    #everything to lower script
    $onlyShort
    $includeShort
    $checkIP

    $servers = $add.split(";")
    foreach ($server in $servers) {
        #check if we want to add local or specific host
        if ($server -match "local") {
            $record = getLocalInfo
        }
        else {
            $record = $server.split("/")
        }
        if ((isIP $record[0])) {
            if ($onlyShort -eq "yes") {
                $short = "only"
                addServer $workingFile $record[0] $record[1] $short
            }
            elseif (($onlyShort -eq "no") -or ($onlyShort -ne "yes")) {
                if (isFQDN $record[1]) {
                    if ($checkIP -eq "yes") {
                        if (checkServer $workingFile $record[0]) {
                            Write-host "Server with IP "$record[0]" is already in hostfile"
                        }
                        else {
                            if (-not ($includeShort)) {
                                $short = "include"
                            }
                            else {
                                $short = "notInclude"
                            }
                            addServer $workingFile $record[0] $record[1] $short
                        }
                    }
                    else {
                        if (-not ($includeShort)) {
                            $short = "include"
                        }
                        else {
                            $short = "notInclude"
                        }
                        addServer $workingFile $record[0] $record[1] $short
                    }
                }
            }
        }
    }
}

#remove servers from hostfile
if ($remove -ne "") {
    $ips = $remove.split(";")
    foreach ($ip in $ips) {
        if ($ip -match "local") {
            $record = getLocalInfo
            $ip = $record[0]
            if(isIP $ip) {
                if(checkServer $workingFile $ip) {
                    removeServer $workingFile $ip
                }
                else {
                    Write-Host "Server with $ip is not in hostfile"
                }
            }

        }
        else {
            if(isIP $ip) {
                if(checkServer $workingFile $ip) {
                    removeServer $workingFile $ip
                }
                else {
                    Write-Host "Server with $ip is not in hostfile"
                }
            }
        }
    }
}
try {
    Copy-Item $workingFile -Destination $originalFile -Force #copying new version of host file to its original location
    Remove-Item $workingFile #removing temporary file
}
catch {
    Write-Host "Unable to replace hosts file. Please try to copy manually. Temp file in C:\Temp\hoststemp.txt."
    exit
}