Update test.ps1

test.ps1

@@ -1,29 +1,57 @@
 function Set-NetworkHardening {
     <#
     .SYNOPSIS
-    Disables legacy or insecure network protocols and services.
+    Hardens the network stack by disabling legacy and insecure protocols and services.
 
     .DESCRIPTION
-    Applies registry and system settings to harden network configuration:
+    This cmdlet applies a set of hardening actions via registry and Windows feature settings.
+    It allows selectively disabling:
+
     - IPv6
-    - LLMNR
+    - LLMNR (Link-Local Multicast Name Resolution)
-    - NBNS
+    - NBNS (NetBIOS over TCP/IP)
-    - mDNS
+    - 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 -DisableIPv6 -DisableLLMNR -DisableNBNS -DisableMDNS -DisableSMB1 -DisableSMB2
+    Set-NetworkHardening -All
+
+    .EXAMPLE
+    Set-NetworkHardening -DisableLLMNR -DisableSMB1
+
+    .EXAMPLE
+    Set-NetworkHardening -Reset
 
     .NOTES
-    Some settings may require a reboot.
+    Author: SVSMSP Toolkit
+    Some changes may require a reboot to take full effect.
     #>
 
     [CmdletBinding()]
@@ -33,9 +61,25 @@ function Set-NetworkHardening {
         [switch]$DisableNBNS,
         [switch]$DisableMDNS,
         [switch]$DisableSMB1,
-        [switch]$DisableSMB2
+        [switch]$DisableSMB2,
+        [switch]$All,
+        [switch]$Reset
     )
 
+    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"
@@ -64,12 +108,11 @@ function Set-NetworkHardening {
         Write-Host "`n[NBNS] Disabling NetBIOS over TCP/IP on all adapters..." -ForegroundColor Cyan
 
         Get-ChildItem -Path $regPath | ForEach-Object {
-            $path = $_.PsPath
             try {
-                Set-ItemProperty -Path $path -Name NetbiosOptions -Value 2
+                Set-ItemProperty -Path $_.PsPath -Name NetbiosOptions -Value 2
-                Write-Host " → Set $path\NetbiosOptions = 2"
+                Write-Host " → Set $($_.PsPath)\NetbiosOptions = 2"
             } catch {
-                Write-Warning "Failed to set NetbiosOptions on $path - $_"
+                Write-Warning "Failed to set NetbiosOptions for $($_.PsPath)"
             }
         }
     }
@@ -109,7 +152,26 @@ function Set-NetworkHardening {
     }
 
     function Disable-SMB2And3 {
-        Write-Host "`n[SMB2/3] Disabling SMB 2.0 and 3.0..." -ForegroundColor Cyan
+        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
@@ -117,7 +179,49 @@ function Set-NetworkHardening {
         Write-Host " → Set $reg\$name = $value"
     }
 
-    # --- Execute requested settings ---
+    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)"
+            }
+        }
+
+        try {
+            Enable-WindowsOptionalFeature -Online -FeatureName SMB1Protocol -NoRestart -ErrorAction Stop
+            Write-Host " → Enabled SMB1 via Windows Feature"
+        } catch {
+            Write-Warning "Could not enable SMB1 via Windows Feature"
+        }
+
+        $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 ($DisableIPv6)  { Disable-IPv6 }
     if ($DisableLLMNR) { Disable-LLMNR }
     if ($DisableNBNS)  { Disable-NBNS }