function Set-NetworkHardening {
    <#
    .SYNOPSIS
    Hardens the network stack by disabling legacy and insecure protocols and services.

    .DESCRIPTION
    This cmdlet applies a set of hardening actions via registry and Windows feature settings.
    It allows selectively disabling:

    - IPv6
    - LLMNR (Link-Local Multicast Name Resolution)
    - NBNS (NetBIOS over TCP/IP)
    - mDNS (Multicast DNS)
    - SMB 1.0
    - SMB 2.0 (also disables SMB 3.0)

    You can use individual parameters or apply all changes using `-All`. Use `-Reset` to undo all changes.

    .PARAMETER DisableIPv6
    Disables IPv6 networking support by modifying the Tcpip6 registry key.

    .PARAMETER DisableLLMNR
    Disables LLMNR name resolution via DNSClient group policy registry.

    .PARAMETER DisableNBNS
    Disables NetBIOS over TCP/IP on all network adapters.

    .PARAMETER DisableMDNS
    Disables mDNS multicast resolution support via dnscache registry key.

    .PARAMETER DisableSMB1
    Disables SMB 1.0 via Windows Features or registry fallback.

    .PARAMETER DisableSMB2
    Disables SMB 2.0 and SMB 3.0 by setting SMB2 = 0 in the LanmanServer registry key (only if no custom shares exist).

    .PARAMETER All
    Enables all above parameters unless explicitly overridden.

    .PARAMETER Reset
    Reverts all applied changes to their default state.

    .EXAMPLE
    Set-NetworkHardening -All

    .EXAMPLE
    Set-NetworkHardening -DisableLLMNR -DisableSMB1

    .EXAMPLE
    Set-NetworkHardening -Reset

    .NOTES
    Author: SVSMSP Toolkit
    Some changes may require a reboot to take full effect.
    #>

    [CmdletBinding()]
    param (
        [switch]$DisableIPv6,
        [switch]$DisableLLMNR,
        [switch]$DisableNBNS,
        [switch]$DisableMDNS,
        [switch]$DisableSMB1,
        [switch]$DisableSMB2,
        [switch]$All,
        [switch]$Reset
    )

    function Reset-NetworkHardening {
        Write-Host "`n[RESET] Reverting all settings to default..." -ForegroundColor Yellow

        $keysToRemove = @(
            @{ Path = "HKLM:\SYSTEM\CurrentControlSet\Services\Tcpip6\Parameters"; Name = "DisabledComponents" },
            @{ Path = "HKLM:\SOFTWARE\Policies\Microsoft\Windows NT\DNSClient"; Name = "EnableMulticast" },
            @{ Path = "HKLM:\SYSTEM\CurrentControlSet\Services\Dnscache\Parameters"; Name = "EnableMDNS" },
            @{ Path = "HKLM:\SYSTEM\CurrentControlSet\Services\LanmanServer\Parameters"; Name = "SMB2" }
        )

        foreach ($key in $keysToRemove) {
            if (Test-Path $key.Path) {
                Remove-ItemProperty -Path $key.Path -Name $key.Name -ErrorAction SilentlyContinue
                Write-Host " → Removed $($key.Path)\$($key.Name)"
            }
        }

        $nbnsPath = "HKLM:\SYSTEM\CurrentControlSet\Services\NetBT\Parameters\Interfaces"
        Get-ChildItem -Path $nbnsPath | ForEach-Object {
            try {
                Set-ItemProperty -Path $_.PsPath -Name NetbiosOptions -Value 0
                Write-Host " → Set $($_.PsPath)\NetbiosOptions = 0"
            } catch {
                Write-Warning "Failed to reset NetbiosOptions for $($_.PsPath)"
            }
        }

        $smb1reg = "HKLM:\SYSTEM\CurrentControlSet\Services\mrxsmb10"
        if (Test-Path $smb1reg) {
            Set-ItemProperty -Path $smb1reg -Name Start -Value 3
            Write-Host " → Set $smb1reg\Start = 3"
        }

        Write-Host "`n✅ Reset complete. Reboot may be required." -ForegroundColor Green
    }

    if ($Reset) {
        Reset-NetworkHardening
        return
    }

    if ($All) {
        if (-not ($DisableIPv6.IsPresent))  { $DisableIPv6  = $true }
        if (-not ($DisableLLMNR.IsPresent)) { $DisableLLMNR = $true }
        if (-not ($DisableNBNS.IsPresent))  { $DisableNBNS  = $true }
        if (-not ($DisableMDNS.IsPresent))  { $DisableMDNS  = $true }
        if (-not ($DisableSMB1.IsPresent))  { $DisableSMB1  = $true }
        if (-not ($DisableSMB2.IsPresent))  { $DisableSMB2  = $true }
    }

    function Disable-IPv6 {
        $regPath = "HKLM:\SYSTEM\CurrentControlSet\Services\Tcpip6\Parameters"
        $name = "DisabledComponents"
        $value = 0xFF

        Write-Host "`n[IPv6] Disabling via registry..." -ForegroundColor Cyan
        New-ItemProperty -Path $regPath -Name $name -PropertyType DWORD -Value $value -Force | Out-Null
        Write-Host " → Set $regPath\$name = $value"
    }

    function Disable-LLMNR {
        $regPath = "HKLM:\SOFTWARE\Policies\Microsoft\Windows NT\DNSClient"
        $name = "EnableMulticast"
        $value = 0

        Write-Host "`n[LLMNR] Disabling via registry..." -ForegroundColor Cyan
        if (-not (Test-Path $regPath)) {
            New-Item -Path $regPath -Force | Out-Null
        }
        New-ItemProperty -Path $regPath -Name $name -PropertyType DWORD -Value $value -Force | Out-Null
        Write-Host " → Set $regPath\$name = $value"
    }

    function Disable-NBNS {
        $regPath = "HKLM:\SYSTEM\CurrentControlSet\Services\NetBT\Parameters\Interfaces"
        Write-Host "`n[NBNS] Disabling NetBIOS over TCP/IP on all adapters..." -ForegroundColor Cyan

        Get-ChildItem -Path $regPath | ForEach-Object {
            try {
                Set-ItemProperty -Path $_.PsPath -Name NetbiosOptions -Value 2
                Write-Host " → Set $($_.PsPath)\NetbiosOptions = 2"
            } catch {
                Write-Warning "Failed to set NetbiosOptions for $($_.PsPath)"
            }
        }
    }

    function Disable-mDNS {
        $regPath = "HKLM:\SYSTEM\CurrentControlSet\Services\Dnscache\Parameters"
        $name = "EnableMDNS"
        $value = 0

        Write-Host "`n[mDNS] Disabling via registry..." -ForegroundColor Cyan
        if (-not (Test-Path $regPath)) {
            New-Item -Path $regPath -Force | Out-Null
        }
        New-ItemProperty -Path $regPath -Name $name -PropertyType DWORD -Value $value -Force | Out-Null
        Write-Host " → Set $regPath\$name = $value"
    }

    function Disable-SMB1 {
        Write-Host "`n[SMB1] Disabling SMB 1.0..." -ForegroundColor Cyan
        try {
            Disable-WindowsOptionalFeature -Online -FeatureName SMB1Protocol -NoRestart -ErrorAction Stop
            Write-Host " → Disabled via Windows Optional Features"
        } catch {
            Write-Warning "SMB1 disable via feature failed. Falling back to registry."
            $reg = "HKLM:\SYSTEM\CurrentControlSet\Services\LanmanServer\Parameters"
            $name = "SMB1"
            $value = 0
            New-ItemProperty -Path $reg -Name $name -PropertyType DWORD -Value $value -Force | Out-Null
            Write-Host " → Set $reg\$name = $value"
        }

        $clientReg = "HKLM:\SYSTEM\CurrentControlSet\Services\mrxsmb10"
        if (Test-Path $clientReg) {
            Set-ItemProperty -Path $clientReg -Name Start -Value 4
            Write-Host " → Set $clientReg\Start = 4 (Disabled)"
        }
    }

    function Disable-SMB2And3 {
        Write-Host "`n[SMB2/3] Checking for custom SMB shares..." -ForegroundColor Cyan

        try {
            $shares = Get-SmbShare | Where-Object {
                $_.Name -notin @('ADMIN$', 'IPC$') -and -not ($_.Name -match '^[A-Z]\$')
            }

            if ($shares.Count -gt 0) {
                Write-Warning "Custom SMB shares detected. SMB 2/3 will NOT be disabled."
                $shares | ForEach-Object {
                    Write-Host " → $($_.Name)  [$($_.Path)]"
                }
                return
            }
        } catch {
            Write-Warning "Failed to query SMB shares: $_"
            return
        }

        Write-Host "No custom SMB shares found. Disabling SMB 2.0 and 3.0..." -ForegroundColor Cyan
        $reg = "HKLM:\SYSTEM\CurrentControlSet\Services\LanmanServer\Parameters"
        $name = "SMB2"
        $value = 0
        New-ItemProperty -Path $reg -Name $name -PropertyType DWORD -Value $value -Force | Out-Null
        Write-Host " → Set $reg\$name = $value"
    }

    if ($DisableIPv6)  { Disable-IPv6 }
    if ($DisableLLMNR) { Disable-LLMNR }
    if ($DisableNBNS)  { Disable-NBNS }
    if ($DisableMDNS)  { Disable-mDNS }
    if ($DisableSMB1)  { Disable-SMB1 }
    if ($DisableSMB2)  { Disable-SMB2And3 }

    Write-Host "`n✅ Network hardening complete. Reboot may be required for changes to apply." -ForegroundColor Green
}