<#
.SYNOPSIS
U盘密钥安全锁系统 - 通过物理U盘验证保护计算机安全

.DESCRIPTION
本脚本提供基于U盘的计算机安全锁定解决方案：
1. 系统启动时强制验证U盘密钥
2. 定期检测U盘密钥状态
3. 未通过验证时自动关机保护

核心功能：
- 启动验证：系统启动后5秒内需插入正确密钥U盘
- 周期验证：每30分钟自动检测密钥状态
- 二次验证：检测失败后给予10分钟宽限期
- 安全关机：验证失败后5秒内强制关机
- 自保护机制：密钥文件自动隐藏+系统属性

配置模式：
powershell -ExecutionPolicy Bypass -File "本脚本路径" -i

卸载命令：
powershell -ExecutionPolicy Bypass -File "本脚本路径" -uninstall

安全特性：
1. 128位高强度随机密钥
2. U盘密钥文件隐藏保护
3. 双计划任务相互独立
4. 配置参数自更新机制
#>

#region 可配置参数
$config = @{
    # 基础配置
    StartupDelay    = 5      # 系统启动后延迟检测时间（秒）
    CheckInterval   = 30     # 计划任务检测间隔（分钟）
    RetryDelay      = 10     # 第二次验证等待时间（分钟）

    # U盘验证配置
    VolumeName      = "Lock_PC"      # 默认卷名
    FileName        = "lock.key"     # 默认文件名
    SecretString    = "kX7@z!9Pq$2*F5%vR1lY#dW8^gH3&sK6(jO0)uC4-bN_mA/eZ+tIxSfLcV=G{B}nM" # 默认密钥（128字符）

    # 计划任务配置
    StartupTaskName = "USBLockMonitor_Startup"   # 启动验证任务名称
    IntervalTaskName = "USBLockMonitor_Interval" # 间隔验证任务名称
    ScriptPath      = $MyInvocation.MyCommand.Path  # 脚本自身路径
}
#endregion

param(
    [switch]$i,           # 配置模式开关
    [switch]$StartupMode, # 启动项运行模式
    [switch]$IntervalMode,# 间隔验证模式
    [switch]$uninstall    # 卸载模式开关
)

# 加载必要的程序集
Add-Type -AssemblyName System.Windows.Forms
Add-Type -AssemblyName System.Security

# 验证U盘密钥函数
function Test-USBKey {
    param(
        [string]$VolumeName,
        [string]$FileName,
        [string]$SecretString
    )

    try {
        # 查找匹配卷名的可移动磁盘
        $targetDrive = Get-Volume | Where-Object {
            $_.DriveType -eq 'Removable' -and
            $_.FileSystemLabel -eq $VolumeName -and
            $_.DriveLetter
        } -ErrorAction Stop

        if (-not $targetDrive) {
            Write-Host "[验证失败] 未找到匹配的U盘: $VolumeName" -ForegroundColor Yellow
            return $false
        }

        # 构建文件路径
        $filePath = "$($targetDrive.DriveLetter):\$FileName"

        # 检查文件是否存在
        if (-not (Test-Path $filePath -ErrorAction Stop)) {
            Write-Host "[验证失败] 未找到密钥文件: $filePath" -ForegroundColor Yellow
            return $false
        }

        # 读取文件内容并验证
        $content = Get-Content $filePath -Raw -ErrorAction Stop
        $result = ($content.Trim() -eq $SecretString)

        if (-not $result) {
            Write-Host "[验证失败] 密钥不匹配" -ForegroundColor Red
        } else {
            Write-Host "[验证成功] 密钥验证通过" -ForegroundColor Green
        }
        return $result
    }
    catch {
        Write-Host "[验证异常] 操作失败: $_" -ForegroundColor Red
        return $false
    }
}

# 显示定时消息函数
function Show-TimedMessage {
    param(
        [string]$Message,
        [int]$TimeoutSeconds,
        [string]$Title = "安全警告",
        [string]$MessageType = "Warning"
    )

    try {
        # 创建独立进程显示弹窗
        $psCommand = @"
        [Console]::OutputEncoding = [Text.Encoding]::UTF8
        Add-Type -AssemblyName System.Windows.Forms
        `$form = New-Object Windows.Forms.Form
        `$form.Text = "$Title"
        `$form.Size = New-Object Drawing.Size(500, 200)
        `$form.StartPosition = "CenterScreen"
        `$form.TopMost = `$true

        `$label = New-Object Windows.Forms.Label
        `$label.Text = "$Message`n`n窗口将在 $TimeoutSeconds 秒后自动关闭"
        `$label.AutoSize = `$true
        `$label.Location = New-Object Drawing.Point(20, 20)
        `$label.Font = New-Object Drawing.Font("Microsoft YaHei", 10)
        `$form.Controls.Add(`$label)

        # 倒计时时钟
        `$timer = New-Object Windows.Forms.Timer
        `$timer.Interval = 1000
        `$count = $TimeoutSeconds

        `$timer.Add_Tick({
            `$count--
            `$label.Text = "$Message`n`n窗口将在 `$count 秒后自动关闭"
            if (`$count -le 0) {
                `$timer.Stop()
                `$form.Close()
            }
        })

        `$timer.Start()
        `$form.ShowDialog()
"@
        Start-Process powershell.exe -ArgumentList "-Command $psCommand" -WindowStyle Normal -ErrorAction Stop
    }
    catch {
        Write-Host "显示消息窗口失败: $_" -ForegroundColor Yellow
        Write-Host "消息内容: $Message" -ForegroundColor Yellow
    }
}

# 关机函数
function Invoke-Shutdown {
    param(
        [switch]$StartupShutdown
    )

    try {
        if ($StartupShutdown) {
            # 启动模式使用控制台警告
            Write-Host "`n[安全警报] 未检测到有效密钥，系统将在5秒后关闭！" -ForegroundColor Red
            Start-Sleep 5
        } else {
            # 其他模式使用GUI警告
            Show-TimedMessage "未插入开机密钥优盘，电脑将在5秒后关闭！" 5
            Start-Sleep 5
        }
        Stop-Computer -Force -ErrorAction Stop
    }
    catch {
        Write-Host "关机失败: $_" -ForegroundColor Red
        Write-Host "请手动保存工作并关闭计算机" -ForegroundColor Yellow
    }
}

# 配置模式函数
function Invoke-Configuration {
    # 请求管理员权限
    if (-not ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
        Write-Host "需要管理员权限进行配置..." -ForegroundColor Yellow
        $scriptPath = $config.ScriptPath
        Start-Process powershell.exe "-NoProfile -ExecutionPolicy Bypass -Command `"& {& '$scriptPath' -i}`"" -Verb RunAs
        exit
    }

    # 保存原始配置用于回滚
    $originalConfig = @{
        VolumeName = $config.VolumeName
        FileName = $config.FileName
        SecretString = $config.SecretString
    }

    try {
        Write-Host "`n配置模式 - U盘密钥安全锁系统" -ForegroundColor Cyan
        Write-Host "请插入用于配置的U盘" -ForegroundColor Cyan
        
        # 获取用户输入
        $driveLetter = Read-Host "`n请输入U盘盘符 (例如 F: )"
        if (-not $driveLetter.EndsWith(':')) {
            $driveLetter += ':'
        }

        # 验证驱动器
        $drive = Get-Volume -DriveLetter $driveLetter[0] -ErrorAction Stop
        if (-not $drive -or $drive.DriveType -ne 'Removable') {
            throw "指定的驱动器不是有效的可移动驱动器!"
        }

        # 获取用户配置
        $newVolume = Read-Host "`n请输入卷名 [$($config.VolumeName)]"
        if ([string]::IsNullOrWhiteSpace($newVolume)) {
            $newVolume = $config.VolumeName
        }

        $newFile = Read-Host "`n请输入文件名 [$($config.FileName)]"
        if ([string]::IsNullOrWhiteSpace($newFile)) {
            $newFile = $config.FileName
        }

        $newSecret = Read-Host "`n请输入密钥 (留空自动生成)"
        if ([string]::IsNullOrWhiteSpace($newSecret)) {
            # 生成128位随机密钥
            $chars = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()_-+=[]{}|;:,.<>?'
            $newSecret = -join (1..128 | ForEach-Object { $chars[(Get-Random -Maximum $chars.Length)] })
            Write-Host "`n已生成新密钥: $newSecret" -ForegroundColor Cyan
        }

        # 更新配置
        $config.VolumeName = $newVolume
        $config.FileName = $newFile
        $config.SecretString = $newSecret

        # 更新U盘
        try {
            # 实时检测U盘连接状态
            $drive = Get-Volume -DriveLetter $driveLetter[0] -ErrorAction Stop
            if (-not $drive -or $drive.DriveType -ne 'Removable') {
                throw "U盘已被移除"
            }

            # 设置卷名
            Set-Volume -DriveLetter $driveLetter[0] -NewFileSystemLabel $newVolume -ErrorAction Stop

            # 实时检测U盘连接状态
            $drive = Get-Volume -DriveLetter $driveLetter[0] -ErrorAction Stop
            if (-not $drive -or $drive.DriveType -ne 'Removable') {
                throw "U盘在设置卷名后被移除"
            }

            # 创建密钥文件
            $keyPath = "$driveLetter\$newFile"
            Set-Content -Path $keyPath -Value $newSecret -Force -ErrorAction Stop

            # 设置文件属性
            attrib +s +h +r $keyPath

            # 再次验证U盘状态
            $drive = Get-Volume -DriveLetter $driveLetter[0] -ErrorAction Stop
            if (-not $drive -or $drive.DriveType -ne 'Removable') {
                throw "U盘在创建密钥文件后被移除"
            }

            Write-Host "`nU盘配置成功!" -ForegroundColor Green
            Write-Host "卷名: $newVolume"
            Write-Host "文件: $newFile"
            Write-Host "密钥: $newSecret"
        }
        catch {
            throw "U盘配置过程中出错: $_"
        }

        # 更新脚本自身配置
        try {
            $scriptContent = Get-Content $config.ScriptPath -Raw -ErrorAction Stop

            $newConfig = @"
#region 可配置参数
`$config = @{
    StartupDelay    = $($config.StartupDelay)
    CheckInterval   = $($config.CheckInterval)
    RetryDelay      = $($config.RetryDelay)
    VolumeName      = "$($config.VolumeName)"
    FileName        = "$($config.FileName)"
    SecretString    = "$($config.SecretString)"
    StartupTaskName = "$($config.StartupTaskName)"
    IntervalTaskName = "$($config.IntervalTaskName)"
    ScriptPath      = `$MyInvocation.MyCommand.Path
}
#endregion
"@

            $newScript = $scriptContent -replace '(?s)#region 可配置参数.*?#endregion', $newConfig
            Set-Content -Path $config.ScriptPath -Value $newScript -Force -ErrorAction Stop

            Write-Host "`n脚本配置已更新!" -ForegroundColor Green
        }
        catch {
            throw "更新脚本配置时出错: $_"
        }

        # 删除可能存在的旧任务
        try {
            $tasks = @($config.StartupTaskName, $config.IntervalTaskName)
            foreach ($task in $tasks) {
                if (Get-ScheduledTask -TaskName $task -ErrorAction SilentlyContinue) {
                    Unregister-ScheduledTask -TaskName $task -Confirm:$false -ErrorAction Stop
                    Write-Host "已删除旧任务: $task" -ForegroundColor Green
                }
            }
        }
        catch {
            Write-Host "删除旧任务时出错: $_" -ForegroundColor Yellow
        }

        # 创建启动验证任务（系统启动时运行）
        try {
            $action = New-ScheduledTaskAction `
                -Execute 'powershell.exe' `
                -Argument "-ExecutionPolicy Bypass -File `"$($config.ScriptPath)`" -StartupMode"

            $trigger = New-ScheduledTaskTrigger `
                -AtStartup `
                -RandomDelay "00:00:$($config.StartupDelay)"

            $settings = New-ScheduledTaskSettingsSet `
                -AllowStartIfOnBatteries `
                -DontStopIfGoingOnBatteries `
                -StartWhenAvailable `
                -ExecutionTimeLimit (New-TimeSpan -Minutes 5)

            $principal = New-ScheduledTaskPrincipal -UserId "SYSTEM" -RunLevel Highest

            Register-ScheduledTask `
                -TaskName $config.StartupTaskName `
                -Action $action `
                -Trigger $trigger `
                -Settings $settings `
                -Principal $principal `
                -Force `
                -ErrorAction Stop | Out-Null

            Write-Host "`n启动验证任务创建成功! 系统启动后 $($config.StartupDelay) 秒检测" -ForegroundColor Green
        }
        catch {
            throw "创建启动验证任务时出错: $_"
        }

        # 创建间隔验证任务（每30分钟运行）
        try {
            $action = New-ScheduledTaskAction `
                -Execute 'powershell.exe' `
                -Argument "-ExecutionPolicy Bypass -File `"$($config.ScriptPath)`" -IntervalMode"

            $trigger = New-ScheduledTaskTrigger `
                -Once `
                -At (Get-Date) `
                -RepetitionInterval (New-TimeSpan -Minutes $config.CheckInterval) `
                -RepetitionDuration ([TimeSpan]::MaxValue)

            $settings = New-ScheduledTaskSettingsSet `
                -AllowStartIfOnBatteries `
                -DontStopIfGoingOnBatteries `
                -StartWhenAvailable `
                -ExecutionTimeLimit (New-TimeSpan -Minutes 30)

            $principal = New-ScheduledTaskPrincipal -UserId "SYSTEM" -RunLevel Highest

            Register-ScheduledTask `
                -TaskName $config.IntervalTaskName `
                -Action $action `
                -Trigger $trigger `
                -Settings $settings `
                -Principal $principal `
                -Force `
                -ErrorAction Stop | Out-Null

            Write-Host "`n间隔验证任务创建成功! 每 $($config.CheckInterval) 分钟检测一次" -ForegroundColor Green
        }
        catch {
            throw "创建间隔验证任务时出错: $_"
        }

        Write-Host "`n配置完成! 按任意键退出..."
        $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    }
    catch {
        # 回滚配置
        $config.VolumeName = $originalConfig.VolumeName
        $config.FileName = $originalConfig.FileName
        $config.SecretString = $originalConfig.SecretString

        Write-Host "`n配置失败: $_" -ForegroundColor Red
        Write-Host "已恢复原始配置" -ForegroundColor Yellow
        Write-Host "按任意键退出..."
        $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
        exit 1
    }
}

# 卸载函数
function Invoke-Uninstall {
    # 请求管理员权限
    if (-not ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
        Write-Host "需要管理员权限进行卸载..." -ForegroundColor Yellow
        $scriptPath = $config.ScriptPath
        Start-Process powershell.exe "-NoProfile -ExecutionPolicy Bypass -Command `"& {& '$scriptPath' -uninstall}`"" -Verb RunAs
        exit
    }

    try {
        Write-Host "`n开始卸载U盘密钥安全锁系统..." -ForegroundColor Cyan

        # 停止所有相关进程
        try {
            $processes = Get-Process -Name powershell -ErrorAction SilentlyContinue | Where-Object {
                $_.CommandLine -like "*$($config.StartupTaskName)*" -or
                $_.CommandLine -like "*$($config.IntervalTaskName)*"
            }

            if ($processes) {
                $processes | Stop-Process -Force
                Write-Host "已停止 $($processes.Count) 个相关进程" -ForegroundColor Green
            } else {
                Write-Host "未找到相关运行进程" -ForegroundColor Yellow
            }
        }
        catch {
            Write-Host "停止进程时出错: $_" -ForegroundColor Red
        }

        # 删除计划任务
        $tasksDeleted = 0
        try {
            $tasks = @($config.StartupTaskName, $config.IntervalTaskName)
            foreach ($task in $tasks) {
                if (Get-ScheduledTask -TaskName $task -ErrorAction SilentlyContinue) {
                    Unregister-ScheduledTask -TaskName $task -Confirm:$false -ErrorAction Stop
                    $tasksDeleted++
                    Write-Host "已删除计划任务: $task" -ForegroundColor Green
                }
            }

            if ($tasksDeleted -eq 0) {
                Write-Host "未找到计划任务" -ForegroundColor Yellow
            }
        }
        catch {
            Write-Host "删除计划任务时出错: $_" -ForegroundColor Red
        }

        # 删除U盘密钥文件
        $drivesCleaned = 0
        try {
            $targetDrives = Get-Volume | Where-Object {
                $_.DriveType -eq 'Removable' -and
                $_.FileSystemLabel -eq $config.VolumeName -and
                $_.DriveLetter
            } -ErrorAction SilentlyContinue

            if ($targetDrives) {
                foreach ($drive in $targetDrives) {
                    try {
                        $filePath = "$($drive.DriveLetter):\$($config.FileName)"
                        if (Test-Path $filePath -ErrorAction SilentlyContinue) {
                            # 取消隐藏属性
                            attrib -s -h -r $filePath
                            Remove-Item $filePath -Force -ErrorAction Stop
                            Write-Host "已删除U盘($($drive.DriveLetter):)上的密钥文件" -ForegroundColor Green
                            $drivesCleaned++
                        }
                    }
                    catch {
                        Write-Host "删除U盘($($drive.DriveLetter):)文件失败: $_" -ForegroundColor Yellow
                    }
                }
            }

            if ($drivesCleaned -eq 0) {
                Write-Host "未找到匹配的U盘密钥文件，请手动删除" -ForegroundColor Yellow
            }
        }
        catch {
            Write-Host "扫描U盘时出错: $_" -ForegroundColor Red
        }

        # 清理日志文件
        $logPath = [System.IO.Path]::Combine([Environment]::GetFolderPath('LocalApplicationData'), "USBLockMonitor")
        if (Test-Path $logPath) {
            try {
                Remove-Item $logPath -Recurse -Force
                Write-Host "已删除日志目录: $logPath" -ForegroundColor Green
            }
            catch {
                Write-Host "删除日志目录时出错: $_" -ForegroundColor Yellow
            }
        }

        # 删除临时文件
        try {
            $tempFiles = Get-ChildItem -Path $env:TEMP -Filter "USBLock_*" -ErrorAction SilentlyContinue
            if ($tempFiles) {
                $tempFiles | Remove-Item -Force -Recurse
                Write-Host "已清理 $($tempFiles.Count) 个临时文件" -ForegroundColor Green
            } else {
                Write-Host "未找到临时文件" -ForegroundColor Yellow
            }
        }
        catch {
            Write-Host "清理临时文件时出错: $_" -ForegroundColor Yellow
        }

        # 询问是否删除脚本自身
        Add-Type -AssemblyName System.Windows.Forms
        $choice = [System.Windows.Forms.MessageBox]::Show(
            "是否要删除此脚本文件？`n`n路径: $($config.ScriptPath)",
            "卸载完成",
            [System.Windows.Forms.MessageBoxButtons]::YesNo,
            [System.Windows.Forms.MessageBoxIcon]::Question
        )

        if ($choice -eq [System.Windows.Forms.DialogResult]::Yes) {
            try {
                # 创建临时批处理文件来删除自身
                $batContent = @"
@echo off
chcp 65001 > nul
echo 正在删除脚本文件...
timeout /t 3 /nobreak >nul
del /f /q "%~s0" >nul
if exist "%~s0" (
    echo 文件删除失败，请手动删除: %~s0
    pause
) else (
    echo 卸载完成!
)
del /f /q "%~f0" >nul
"@
                $batPath = [System.IO.Path]::GetTempFileName() + ".bat"
                Set-Content -Path $batPath -Value $batContent -Encoding UTF8

                # 启动批处理
                Start-Process $batPath -WindowStyle Hidden
                Write-Host "脚本将在3秒后自删除" -ForegroundColor Magenta
                exit
            }
            catch {
                Write-Host "删除脚本失败，请手动删除: $($config.ScriptPath)" -ForegroundColor Red
                pause
            }
        } else {
            [System.Windows.Forms.MessageBox]::Show(
                "卸载完成! 脚本文件保留在原始位置。",
                "系统提示",
                [System.Windows.Forms.MessageBoxButtons]::OK,
                [System.Windows.Forms.MessageBoxIcon]::Information
            )
            Write-Host "卸载完成! 脚本文件保留在原始位置。" -ForegroundColor Green
        }
    }
    catch {
        Write-Host "卸载过程中发生严重错误: $_" -ForegroundColor Red
        Write-Host "请手动完成卸载操作" -ForegroundColor Yellow
        pause
        exit 1
    }
}

# 启动守护进程函数
function Start-GuardedProcess {
    param(
        [string]$Command
    )

    try {
        $guardScript = @"
        Start-Process powershell.exe -ArgumentList "$Command" -WindowStyle Hidden -Wait
"@

        Start-Process powershell.exe -ArgumentList "-Command $guardScript" -WindowStyle Hidden -ErrorAction Stop
    }
    catch {
        Write-Host "启动守护进程失败: $_" -ForegroundColor Yellow
        Write-Host "将直接执行验证命令" -ForegroundColor Yellow

        # 直接执行命令作为备用方案
        Invoke-Expression $Command
    }
}

# 主执行流程
if ($i) {
    Invoke-Configuration
    exit
}

if ($uninstall) {
    Invoke-Uninstall
    exit
}

# 启动项验证模式（系统启动后运行）
if ($StartupMode) {
    try {
        Start-Sleep $config.StartupDelay

        if (-not (Test-USBKey $config.VolumeName $config.FileName $config.SecretString)) {
            Invoke-Shutdown -StartupShutdown
        }
    }
    catch {
        Write-Host "启动验证过程中出错: $_" -ForegroundColor Red
        Invoke-Shutdown -StartupShutdown
    }
    exit
}

# 间隔验证模式（计划任务定期运行）
if ($IntervalMode) {
    try {
        # 第一次验证
        if (Test-USBKey $config.VolumeName $config.FileName $config.SecretString) {
            exit
        }

        # 第一次验证失败，显示警告
        $message = "未插入开机密钥优盘，电脑将在$($config.RetryDelay)分钟后再次检测！`n请保存工作进度，并插入开机密钥优盘；`n否则电脑将关闭！"
        Show-TimedMessage $message 30 "安全警告"

        # 构建验证命令
        $command = @"
        Start-Sleep ($($config.RetryDelay * 60))
        if (-not (Test-USBKey '$($config.VolumeName)' '$($config.FileName)' '$($config.SecretString)')) {
            Invoke-Shutdown
        }
"@

        # 启动守护进程执行等待和第二次验证
        Start-GuardedProcess -Command $command
    }
    catch {
        Write-Host "间隔验证过程中出错: $_" -ForegroundColor Red
        Invoke-Shutdown
    }
}