#Requires -RunAsAdministrator
<#
.SYNOPSIS
    Master installation script for SolidWorks Privacy Lockdown.

.DESCRIPTION
    Runs all privacy lockdown components in sequence:
    1. Blocks telemetry domains via hosts file
    2. Disables update/background services
    3. Configures firewall rules
    4. Disables in-app telemetry via registry
    5. Verifies the lockdown

    Generates a detailed report of all changes made.

.NOTES
    Author: Atomaste Solution
    Requires: Administrator privileges
    Run this script to apply all privacy protections at once.

.EXAMPLE
    .\00-install-privacy-lockdown.ps1
    Applies all privacy protections.

.EXAMPLE
    .\00-install-privacy-lockdown.ps1 -Undo
    Removes all privacy protections and restores original settings.

.EXAMPLE
    .\00-install-privacy-lockdown.ps1 -DryRun
    Shows what would be changed without making any modifications.
#>

param(
    [switch]$Undo,           # Remove all protections
    [switch]$DryRun,         # Preview changes without applying
    [switch]$SkipVerify,     # Skip verification step
    [string]$InstallPath = "C:\Program Files\SOLIDWORKS Corp",
    [string]$ReportPath      # Custom report path (default: script directory)
)

$scriptDir = $PSScriptRoot
$timestamp = Get-Date -Format "yyyy-MM-dd_HH-mm-ss"
$reportDir = if ($ReportPath) { $ReportPath } else { Join-Path $scriptDir "reports" }
$reportFile = Join-Path $reportDir "lockdown-report_$timestamp.txt"

# Global report content
$script:report = @()
$script:changes = @{
    HostsFile = @()
    Services = @()
    Firewall = @()
    Registry = @()
}

function Initialize-Report {
    # Create reports directory if needed
    if (-not (Test-Path $reportDir)) {
        New-Item -Path $reportDir -ItemType Directory -Force | Out-Null
    }

    $header = @"
================================================================================
  SOLIDWORKS PRIVACY LOCKDOWN REPORT
  Generated: $(Get-Date -Format "yyyy-MM-dd HH:mm:ss")
  Computer: $env:COMPUTERNAME
  User: $env:USERNAME
  Mode: $(if ($Undo) { "UNINSTALL" } elseif ($DryRun) { "DRY RUN (preview)" } else { "INSTALL" })
================================================================================

"@
    $script:report += $header
    Write-Host $header -ForegroundColor Cyan
}

function Add-ReportSection {
    param(
        [string]$Title,
        [string]$Content
    )

    $section = @"

--------------------------------------------------------------------------------
  $Title
--------------------------------------------------------------------------------
$Content
"@
    $script:report += $section
}

function Add-ReportLine {
    param([string]$Line)
    $script:report += $Line
}

function Write-Log {
    param(
        [string]$Message,
        [string]$Level = "INFO",  # INFO, OK, WARN, ERROR, CHANGE
        [string]$Category         # HostsFile, Services, Firewall, Registry
    )

    $prefix = switch ($Level) {
        "OK"     { "[OK]    "; $color = "Green" }
        "WARN"   { "[WARN]  "; $color = "Yellow" }
        "ERROR"  { "[ERROR] "; $color = "Red" }
        "CHANGE" { "[CHANGE]"; $color = "Magenta" }
        "SKIP"   { "[SKIP]  "; $color = "Gray" }
        "BEFORE" { "[BEFORE]"; $color = "DarkGray" }
        "AFTER"  { "[AFTER] "; $color = "White" }
        default  { "[INFO]  "; $color = "White" }
    }

    $logLine = "$prefix $Message"
    Write-Host $logLine -ForegroundColor $color
    Add-ReportLine $logLine

    # Track changes by category
    if ($Category -and $Level -eq "CHANGE") {
        $script:changes[$Category] += $Message
    }
}

function Save-Report {
    # Add summary section
    $summary = @"

================================================================================
  CHANGES SUMMARY
================================================================================

HOSTS FILE CHANGES ($($script:changes.HostsFile.Count)):
"@
    if ($script:changes.HostsFile.Count -eq 0) {
        $summary += "`n  (no changes)"
    } else {
        foreach ($change in $script:changes.HostsFile) {
            $summary += "`n  - $change"
        }
    }

    $summary += @"

SERVICES CHANGES ($($script:changes.Services.Count)):
"@
    if ($script:changes.Services.Count -eq 0) {
        $summary += "`n  (no changes)"
    } else {
        foreach ($change in $script:changes.Services) {
            $summary += "`n  - $change"
        }
    }

    $summary += @"

FIREWALL CHANGES ($($script:changes.Firewall.Count)):
"@
    if ($script:changes.Firewall.Count -eq 0) {
        $summary += "`n  (no changes)"
    } else {
        foreach ($change in $script:changes.Firewall) {
            $summary += "`n  - $change"
        }
    }

    $summary += @"

REGISTRY CHANGES ($($script:changes.Registry.Count)):
"@
    if ($script:changes.Registry.Count -eq 0) {
        $summary += "`n  (no changes)"
    } else {
        foreach ($change in $script:changes.Registry) {
            $summary += "`n  - $change"
        }
    }

    $totalChanges = $script:changes.HostsFile.Count + $script:changes.Services.Count +
                    $script:changes.Firewall.Count + $script:changes.Registry.Count

    $summary += @"

================================================================================
  TOTAL CHANGES: $totalChanges
  REPORT SAVED: $reportFile
================================================================================
"@

    $script:report += $summary

    # Write report to file
    $script:report -join "`n" | Set-Content -Path $reportFile -Force -Encoding UTF8

    Write-Host ""
    Write-Host "========================================" -ForegroundColor Green
    Write-Host "  REPORT SAVED" -ForegroundColor Green
    Write-Host "========================================" -ForegroundColor Green
    Write-Host "  Location: $reportFile" -ForegroundColor White
    Write-Host "  Total changes: $totalChanges" -ForegroundColor White
    Write-Host ""
}

function Test-IsAdmin {
    $currentUser = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = New-Object Security.Principal.WindowsPrincipal($currentUser)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Get-SystemState {
    Add-ReportSection -Title "SYSTEM STATE BEFORE CHANGES" -Content ""

    # Hosts file state
    Write-Log "Checking hosts file..." -Level "INFO"
    $hostsPath = "C:\Windows\System32\drivers\etc\hosts"
    $hostsContent = Get-Content $hostsPath -Raw -ErrorAction SilentlyContinue

    if ($hostsContent -match "SOLIDWORKS TELEMETRY BLOCK") {
        Write-Log "Hosts file: SolidWorks block already present" -Level "BEFORE"
    } else {
        Write-Log "Hosts file: No SolidWorks block found" -Level "BEFORE"
    }

    # Services state
    Write-Log "Checking services..." -Level "INFO"
    $services = Get-Service -DisplayName "*SOLIDWORKS*","*3DEXPERIENCE*","*Dassault*" -ErrorAction SilentlyContinue

    if ($services) {
        foreach ($svc in $services) {
            Write-Log "Service: $($svc.DisplayName) - Status: $($svc.Status), StartType: $($svc.StartType)" -Level "BEFORE"
        }
    } else {
        Write-Log "Services: No SolidWorks services found" -Level "BEFORE"
    }

    # Firewall rules state
    Write-Log "Checking firewall rules..." -Level "INFO"
    $rules = Get-NetFirewallRule -DisplayName "SolidWorks Privacy -*" -ErrorAction SilentlyContinue

    if ($rules) {
        foreach ($rule in $rules) {
            Write-Log "Firewall: $($rule.DisplayName) - Action: $($rule.Action), Enabled: $($rule.Enabled)" -Level "BEFORE"
        }
    } else {
        Write-Log "Firewall: No SolidWorks Privacy rules found" -Level "BEFORE"
    }

    # Registry state
    Write-Log "Checking registry..." -Level "INFO"
    $swBasePath = "HKCU:\Software\SolidWorks"

    if (Test-Path $swBasePath) {
        $versions = Get-ChildItem -Path $swBasePath -ErrorAction SilentlyContinue |
                    Where-Object { $_.PSChildName -like "SOLIDWORKS *" }

        foreach ($version in $versions) {
            Write-Log "Registry: Found $($version.PSChildName)" -Level "BEFORE"

            $perfPath = Join-Path $version.PSPath "Performance"
            if (Test-Path $perfPath) {
                $props = Get-ItemProperty -Path $perfPath -ErrorAction SilentlyContinue
                if ($props.PSObject.Properties.Name -contains "OptInStatus") {
                    Write-Log "  CEIP OptInStatus = $($props.OptInStatus)" -Level "BEFORE"
                }
            }
        }
    } else {
        Write-Log "Registry: No SolidWorks keys found (SW may not be installed)" -Level "BEFORE"
    }

    Add-ReportLine ""
}

function Apply-HostsFileBlock {
    Add-ReportSection -Title "HOSTS FILE MODIFICATIONS" -Content ""

    $hostsPath = "C:\Windows\System32\drivers\etc\hosts"
    $backupPath = "C:\Windows\System32\drivers\etc\hosts.backup.solidworks.$timestamp"

    $blockDomains = @(
        "api.3ds.com",
        "www.3ds.com",
        "swym.3ds.com",
        "iam.3ds.com",
        "cas.3ds.com",
        "eu1-ds-iam.3dexperience.3ds.com",
        "eu1-ds.3dexperience.3ds.com",
        "update.solidworks.com",
        "www.solidworks.com",
        "sentry.io",
        "o136956.ingest.sentry.io",
        "telemetry.solidworks.com",
        "analytics.3ds.com",
        "collect.3ds.com",
        "ifwe.3ds.com",
        "eu1-ifwe.3dexperience.3ds.com",
        "passport.3ds.com",
        "3dswym.3ds.com"
    )

    $licensingDomains = @(
        "activation.solidworks.com",
        "license.solidworks.com",
        "licensing.solidworks.com"
    )

    Write-Log "Licensing domains (preserved):" -Level "INFO"
    foreach ($domain in $licensingDomains) {
        Write-Log "  ALLOW: $domain" -Level "OK"
    }

    Write-Log "" -Level "INFO"
    Write-Log "Telemetry domains to block:" -Level "INFO"

    if ($DryRun) {
        foreach ($domain in $blockDomains) {
            Write-Log "  Would block: $domain" -Level "INFO"
        }
        Write-Log "[DRY RUN] No changes made to hosts file" -Level "WARN"
        return
    }

    # Backup
    Copy-Item -Path $hostsPath -Destination $backupPath -Force
    Write-Log "Backup created: $backupPath" -Level "OK"

    $content = Get-Content $hostsPath -Raw
    $markerStart = "# === SOLIDWORKS TELEMETRY BLOCK START ==="
    $markerEnd = "# === SOLIDWORKS TELEMETRY BLOCK END ==="

    # Remove existing block if present
    if ($content -match [regex]::Escape($markerStart)) {
        $pattern = "$([regex]::Escape($markerStart))[\s\S]*?$([regex]::Escape($markerEnd))\r?\n?"
        $content = $content -replace $pattern, ""
        Write-Log "Removed existing block" -Level "INFO"
    }

    # Build new block
    $blockContent = @()
    $blockContent += ""
    $blockContent += $markerStart
    $blockContent += "# Blocks telemetry/analytics while preserving licensing"
    $blockContent += "# Generated: $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')"
    $blockContent += "#"

    foreach ($domain in $blockDomains) {
        $blockContent += "127.0.0.1 $domain"
        Write-Log "BLOCKED: $domain -> 127.0.0.1" -Level "CHANGE" -Category "HostsFile"
    }

    $blockContent += $markerEnd

    # Apply
    Set-Content -Path $hostsPath -Value $content.TrimEnd() -Force -NoNewline
    Add-Content -Path $hostsPath -Value ($blockContent -join "`r`n") -Force

    Write-Log "Hosts file updated successfully" -Level "OK"
}

function Remove-HostsFileBlock {
    Add-ReportSection -Title "HOSTS FILE RESTORATION" -Content ""

    $hostsPath = "C:\Windows\System32\drivers\etc\hosts"
    $content = Get-Content $hostsPath -Raw
    $markerStart = "# === SOLIDWORKS TELEMETRY BLOCK START ==="
    $markerEnd = "# === SOLIDWORKS TELEMETRY BLOCK END ==="

    if ($DryRun) {
        if ($content -match [regex]::Escape($markerStart)) {
            Write-Log "[DRY RUN] Would remove SolidWorks block from hosts file" -Level "WARN"
        } else {
            Write-Log "[DRY RUN] No SolidWorks block found to remove" -Level "INFO"
        }
        return
    }

    if ($content -match [regex]::Escape($markerStart)) {
        $pattern = "$([regex]::Escape($markerStart))[\s\S]*?$([regex]::Escape($markerEnd))\r?\n?"
        $newContent = $content -replace $pattern, ""
        Set-Content -Path $hostsPath -Value $newContent.TrimEnd() -Force -NoNewline
        Add-Content -Path $hostsPath -Value ""
        Write-Log "Removed SolidWorks telemetry block from hosts file" -Level "CHANGE" -Category "HostsFile"
    } else {
        Write-Log "No SolidWorks block found in hosts file" -Level "SKIP"
    }
}

function Manage-Services {
    param([switch]$Restore)

    $title = if ($Restore) { "SERVICES RESTORATION" } else { "SERVICES MODIFICATIONS" }
    Add-ReportSection -Title $title -Content ""

    $stateFile = "$env:USERPROFILE\solidworks-services-backup.json"

    $services = Get-Service -DisplayName "*SOLIDWORKS*","*3DEXPERIENCE*","*Dassault*" -ErrorAction SilentlyContinue

    if (-not $services -or $services.Count -eq 0) {
        Write-Log "No SolidWorks services found on this system" -Level "INFO"
        return
    }

    if ($Restore) {
        if ($DryRun) {
            Write-Log "[DRY RUN] Would restore services from: $stateFile" -Level "WARN"
            return
        }

        if (Test-Path $stateFile) {
            $states = Get-Content -Path $stateFile -Raw | ConvertFrom-Json
            foreach ($serviceName in $states.PSObject.Properties.Name) {
                $originalState = $states.$serviceName
                try {
                    Set-Service -Name $serviceName -StartupType $originalState.StartType -ErrorAction Stop
                    Write-Log "Restored: $($originalState.DisplayName) -> $($originalState.StartType)" -Level "CHANGE" -Category "Services"
                } catch {
                    Write-Log "Failed to restore $serviceName : $_" -Level "ERROR"
                }
            }
        } else {
            Write-Log "No backup file found: $stateFile" -Level "WARN"
        }
        return
    }

    # Save current states
    $states = @{}
    foreach ($svc in $services) {
        $states[$svc.Name] = @{
            DisplayName = $svc.DisplayName
            Status = $svc.Status.ToString()
            StartType = $svc.StartType.ToString()
        }
    }

    if (-not $DryRun) {
        $states | ConvertTo-Json -Depth 3 | Set-Content -Path $stateFile -Force
        Write-Log "Service states backed up to: $stateFile" -Level "OK"
    }

    foreach ($svc in $services) {
        $isUpdateService = $svc.DisplayName -match "update|download|background|3dexperience"

        if ($isUpdateService) {
            if ($DryRun) {
                Write-Log "[DRY RUN] Would disable: $($svc.DisplayName) (currently: $($svc.StartType))" -Level "WARN"
            } else {
                try {
                    if ($svc.Status -eq "Running") {
                        Stop-Service -Name $svc.Name -Force -ErrorAction Stop
                        Write-Log "Stopped: $($svc.DisplayName)" -Level "OK"
                    }
                    Set-Service -Name $svc.Name -StartupType Disabled -ErrorAction Stop
                    Write-Log "Disabled: $($svc.DisplayName) (was: $($svc.StartType))" -Level "CHANGE" -Category "Services"
                } catch {
                    Write-Log "Failed to disable $($svc.DisplayName): $_" -Level "ERROR"
                }
            }
        } else {
            Write-Log "Preserved: $($svc.DisplayName) ($($svc.StartType))" -Level "SKIP"
        }
    }
}

function Manage-Firewall {
    param([switch]$Remove)

    $title = if ($Remove) { "FIREWALL RULES REMOVAL" } else { "FIREWALL RULES CONFIGURATION" }
    Add-ReportSection -Title $title -Content ""

    $rulePrefix = "SolidWorks Privacy - "

    if ($Remove) {
        $rules = Get-NetFirewallRule -DisplayName "$rulePrefix*" -ErrorAction SilentlyContinue

        if ($DryRun) {
            if ($rules) {
                Write-Log "[DRY RUN] Would remove $($rules.Count) firewall rule(s)" -Level "WARN"
            } else {
                Write-Log "[DRY RUN] No firewall rules to remove" -Level "INFO"
            }
            return
        }

        if ($rules) {
            foreach ($rule in $rules) {
                $rules | Remove-NetFirewallRule
                Write-Log "Removed rule: $($rule.DisplayName)" -Level "CHANGE" -Category "Firewall"
            }
        } else {
            Write-Log "No SolidWorks firewall rules found" -Level "SKIP"
        }
        return
    }

    # Find executables to block
    $execsToBlock = @(
        "SOLIDWORKS\swScheduler\dxttasks.exe",
        "SOLIDWORKS\swScheduler\DXTTask.exe",
        "SOLIDWORKS Installation Manager\sldIM.exe",
        "SOLIDWORKS Installation Manager\sldDownloader.exe"
    )

    if (-not (Test-Path $InstallPath)) {
        Write-Log "SolidWorks installation not found at: $InstallPath" -Level "WARN"
        Write-Log "Firewall rules will be created when SolidWorks is installed" -Level "INFO"
        return
    }

    foreach ($exe in $execsToBlock) {
        $fullPath = Join-Path $InstallPath $exe
        $exeName = Split-Path $fullPath -Leaf
        $ruleName = "$rulePrefix Block $exeName"

        if (-not (Test-Path $fullPath)) {
            Write-Log "Executable not found: $exe" -Level "SKIP"
            continue
        }

        $existing = Get-NetFirewallRule -DisplayName $ruleName -ErrorAction SilentlyContinue

        if ($existing) {
            Write-Log "Rule already exists: $exeName" -Level "SKIP"
            continue
        }

        if ($DryRun) {
            Write-Log "[DRY RUN] Would create block rule for: $exeName" -Level "WARN"
        } else {
            try {
                New-NetFirewallRule `
                    -DisplayName $ruleName `
                    -Description "Blocks outbound for SolidWorks update/telemetry" `
                    -Direction Outbound `
                    -Action Block `
                    -Program $fullPath `
                    -Enabled True `
                    -Profile Any `
                    -ErrorAction Stop | Out-Null

                Write-Log "Created block rule: $exeName" -Level "CHANGE" -Category "Firewall"
            } catch {
                Write-Log "Failed to create rule for $exeName : $_" -Level "ERROR"
            }
        }
    }
}

function Manage-Registry {
    param([switch]$Restore)

    $title = if ($Restore) { "REGISTRY RESTORATION" } else { "REGISTRY MODIFICATIONS" }
    Add-ReportSection -Title $title -Content ""

    $swBasePath = "HKCU:\Software\SolidWorks"

    if (-not (Test-Path $swBasePath)) {
        Write-Log "SolidWorks registry keys not found" -Level "WARN"
        Write-Log "Run this script again after launching SolidWorks" -Level "INFO"
        return
    }

    $versions = Get-ChildItem -Path $swBasePath -ErrorAction SilentlyContinue |
                Where-Object { $_.PSChildName -like "SOLIDWORKS *" }

    if (-not $versions) {
        Write-Log "No SolidWorks version keys found" -Level "WARN"
        return
    }

    $settings = @(
        @{ Path = "Performance"; Name = "OptInStatus"; DisableValue = 0; Desc = "CEIP" },
        @{ Path = "General"; Name = "Auto Check for Updates"; DisableValue = 0; Desc = "Auto Updates" },
        @{ Path = "Performance"; Name = "EnableAnalytics"; DisableValue = 0; Desc = "Analytics" },
        @{ Path = "Performance"; Name = "EnableTelemetry"; DisableValue = 0; Desc = "Telemetry" }
    )

    foreach ($version in $versions) {
        Write-Log "Processing: $($version.PSChildName)" -Level "INFO"

        foreach ($setting in $settings) {
            $regPath = Join-Path $version.PSPath $setting.Path

            if (-not (Test-Path $regPath)) {
                Write-Log "  Path not found: $($setting.Path)" -Level "SKIP"
                continue
            }

            $props = Get-ItemProperty -Path $regPath -ErrorAction SilentlyContinue
            $currentValue = $props.PSObject.Properties[$setting.Name].Value

            if ($Restore) {
                # Restore to enabled (1)
                if ($DryRun) {
                    Write-Log "[DRY RUN] Would restore $($setting.Desc) to enabled" -Level "WARN"
                } else {
                    try {
                        Set-ItemProperty -Path $regPath -Name $setting.Name -Value 1 -Type DWord -Force
                        Write-Log "Restored: $($setting.Desc) = 1 (enabled)" -Level "CHANGE" -Category "Registry"
                    } catch {
                        Write-Log "Failed to restore $($setting.Desc): $_" -Level "ERROR"
                    }
                }
            } else {
                # Disable
                if ($null -ne $currentValue -and $currentValue -eq $setting.DisableValue) {
                    Write-Log "  $($setting.Desc): Already disabled" -Level "SKIP"
                    continue
                }

                if ($DryRun) {
                    Write-Log "[DRY RUN] Would disable $($setting.Desc) (current: $currentValue)" -Level "WARN"
                } else {
                    try {
                        Set-ItemProperty -Path $regPath -Name $setting.Name -Value $setting.DisableValue -Type DWord -Force
                        Write-Log "Disabled: $($setting.Desc) (was: $currentValue -> now: $($setting.DisableValue))" -Level "CHANGE" -Category "Registry"
                    } catch {
                        Write-Log "Failed to disable $($setting.Desc): $_" -Level "ERROR"
                    }
                }
            }
        }
    }
}

# Main execution
if (-not (Test-IsAdmin)) {
    Write-Host "[ERROR] This script requires Administrator privileges." -ForegroundColor Red
    Write-Host "Please right-click PowerShell and select 'Run as Administrator'" -ForegroundColor Yellow
    exit 1
}

Initialize-Report

if ($DryRun) {
    Write-Host ""
    Write-Host "*** DRY RUN MODE - No changes will be made ***" -ForegroundColor Yellow
    Write-Host ""
}

# Capture initial state
Get-SystemState

if ($Undo) {
    Add-ReportSection -Title "REMOVING PRIVACY LOCKDOWN" -Content ""

    if (-not $DryRun) {
        Write-Host ""
        $confirm = Read-Host "Are you sure you want to remove privacy protections? (y/N)"
        if ($confirm -ne "y" -and $confirm -ne "Y") {
            Write-Log "Operation cancelled by user" -Level "WARN"
            Save-Report
            exit 0
        }
    }

    Manage-Registry -Restore
    Manage-Firewall -Remove
    Manage-Services -Restore
    Remove-HostsFileBlock

} else {
    Add-ReportSection -Title "APPLYING PRIVACY LOCKDOWN" -Content ""

    if (-not $DryRun) {
        Write-Host ""
        Write-Host "This will configure:" -ForegroundColor White
        Write-Host "  - Block telemetry domains in hosts file" -ForegroundColor Gray
        Write-Host "  - Disable update/background services" -ForegroundColor Gray
        Write-Host "  - Create firewall block rules" -ForegroundColor Gray
        Write-Host "  - Disable telemetry in registry" -ForegroundColor Gray
        Write-Host ""
        Write-Host "Licensing will remain functional." -ForegroundColor Green
        Write-Host ""

        $confirm = Read-Host "Proceed? (Y/n)"
        if ($confirm -eq "n" -or $confirm -eq "N") {
            Write-Log "Operation cancelled by user" -Level "WARN"
            Save-Report
            exit 0
        }
    }

    Apply-HostsFileBlock
    Manage-Services
    Manage-Firewall
    Manage-Registry
}

Save-Report

if ($DryRun) {
    Write-Host ""
    Write-Host "*** DRY RUN COMPLETE - No changes were made ***" -ForegroundColor Yellow
    Write-Host "Review the report above. Run without -DryRun to apply changes." -ForegroundColor Yellow
}