import win32com.client

def create_update_session():
    """創(chuàng)建 Windows Update Session"""
    return win32com.client.Dispatch("Microsoft.Update.Session")

def scan_for_updates():
    """
    掃描系統(tǒng)中可用的更新
    返回: (待安裝列表, 可選安裝列表)
    """
    session = create_update_session()
    searcher = session.CreateUpdateSearcher()

    print("正在掃描可用更新...（可能需要幾分鐘）")

    try:
        # 搜索所有可用更新
        result = searcher.Search("IsInstalled=0")

        pending = []   # 重要/推薦更新
        optional = []  # 可選更新

        for update in result.Updates:
            info = {
                "title": update.Title,
                "description": update.Description[:100] if update.Description else "",
                "kb": _extract_kb(update.Title),
                "severity": _get_severity(update),
                "size_mb": round(update.DownloadContents.TotalSize / 1024 / 1024, 1)
                            if update.DownloadContents.TotalSize > 0 else 0,
                "is_installed": update.IsInstalled,
            }

            if info["severity"] in ("Critical", "Important"):
                pending.append(info)
            else:
                optional.append(info)

        print(f"掃描完成！")
        print(f"  待安裝（重要）: {len(pending)} 個")
        print(f"  可選更新: {len(optional)} 個")
        print(f"  總大小: {sum(u['size_mb'] for u in pending + optional):.1f} MB")

        return pending, optional

    except Exception as e:
        print(f"掃描失敗: {e}")
        return [], []


def _extract_kb(title):
    """從標(biāo)題中提取 KB 編號"""
    import re
    match = re.search(r"(KB\d+)", title, re.IGNORECASE)
    return match.group(1) if match else ""


def _get_severity(update):
    """獲取更新嚴重級別"""
    try:
        for category in update.Categories:
            name = category.Name.lower()
            if "critical" in name:
                return "Critical"
            if "important" in name:
                return "Important"
            if "moderate" in name:
                return "Moderate"
        return "Low"
    except Exception:
        return "Unknown"


# 使用
pending, optional = scan_for_updates()
print("\n重要更新列表：")
for u in pending[:10]:
    print(f"  [{u['severity']}] {u['title']}")
    print(f"    KB: {u['kb']}, 大小: {u['size_mb']} MB")

import time

def download_updates(updates_info, progress_callback=None):
    """
    下載指定的更新
    updates_info: 從 scan_for_updates 獲取的列表
    """
    session = create_update_session()
    downloader = session.CreateUpdateDownloader()

    # 根據(jù)標(biāo)題創(chuàng)建更新集合
    searcher = session.CreateUpdateSearcher()
    result = searcher.Search("IsInstalled=0")

    updates_to_download = win32com.client.Dispatch("Microsoft.Update.UpdateColl")
    for update in result.Updates:
        for info in updates_info:
            if update.Title == info["title"]:
                updates_to_download.Add(update)
                break

    if updates_to_download.Count == 0:
        print("沒有需要下載的更新")
        return True

    downloader.Updates = updates_to_download

    print(f"開始下載 {updates_to_download.Count} 個更新...")

    try:
        download_result = downloader.Download()
        if download_result.ResultCode == 2:  # Succeeded
            print("下載完成！")
            return True
        else:
            print(f"下載失敗，結(jié)果碼: {download_result.ResultCode}")
            return False
    except Exception as e:
        print(f"下載出錯: {e}")
        return False


def install_updates(updates_info, progress_callback=None):
    """
    安裝已下載的更新
    注意：安裝后可能需要重啟
    """
    session = create_update_session()
    installer = session.CreateUpdateInstaller()

    # 獲取已下載的更新
    searcher = session.CreateUpdateSearcher()
    result = searcher.Search("IsInstalled=0 AND IsDownloaded=1")

    updates_to_install = win32com.client.Dispatch("Microsoft.Update.UpdateColl")
    for update in result.Updates:
        for info in updates_info:
            if update.Title == info["title"]:
                updates_to_install.Add(update)
                break

    if updates_to_install.Count == 0:
        print("沒有已下載的更新可供安裝")
        return {"installed": 0, "failed": 0, "reboot_required": False}

    installer.Updates = updates_to_install

    print(f"開始安裝 {updates_to_install.Count} 個更新...")
    print("?? 此操作需要管理員權(quán)限！")

    try:
        install_result = installer.Install()

        installed = install_result.ResultCode
        reboot_required = installer.RebootRequired

        # 統(tǒng)計結(jié)果
        success_count = 0
        fail_count = 0

        for i in range(install_result.Updates.Count):
            result_code = install_result.GetUpdateResult(i).ResultCode
            update_title = install_result.Updates.Item(i).Title
            if result_code == 2:  # Succeeded
                success_count += 1
                print(f"  ? {update_title}")
            else:
                fail_count += 1
                print(f"  ? {update_title} (失敗: {result_code})")

        print(f"\n安裝完成: {success_count} 成功, {fail_count} 失敗")
        if reboot_required:
            print("?? 需要重啟才能完成安裝")

        return {
            "installed": success_count,
            "failed": fail_count,
            "reboot_required": reboot_required,
        }

    except Exception as e:
        print(f"安裝出錯: {e}")
        return {"installed": 0, "failed": 0, "reboot_required": False, "error": str(e)}


# 使用示例（請以管理員身份運行）
# pending, _ = scan_for_updates()
# download_updates(pending)
# install_updates(pending)

def get_installed_updates(days=30):
    """
    獲取最近安裝的更新歷史
    days: 查看最近多少天的安裝記錄
    """
    session = create_update_session()
    searcher = session.CreateUpdateHistory()

    from datetime import datetime, timedelta
    cutoff = datetime.now() - timedelta(days=days)

    history = []
    for entry in searcher:
        try:
            date_str = entry.Date
            install_date = datetime.strptime(
                str(date_str)[:10], "%Y-%m-%d"
            ) if date_str else None

            if install_date and install_date >= cutoff:
                history.append({
                    "title": entry.Title,
                    "date": install_date,
                    "result": "成功" if entry.ResultCode == 2 else "失敗",
                    "kb": _extract_kb(entry.Title),
                })
        except Exception:
            continue

    # 按日期排序
    history.sort(key=lambda x: x["date"], reverse=True)

    print(f"最近 {days} 天安裝的更新 ({len(history)} 個)：\n")
    for h in history[:20]:
        print(f"  [{h['date'].strftime('%Y-%m-%d')}] {h['result']} "
              f"{h['title']}")

    return history


# 使用
# get_installed_updates(days=30)

import subprocess
import json

def scan_updates_powershell():
    """通過 PowerShell 掃描可用更新"""
    ps_cmd = '''
    $UpdateSession = New-Object -ComObject Microsoft.Update.Session
    $UpdateSearcher = $UpdateSession.CreateUpdateSearcher()
    $SearchResult = $UpdateSearcher.Search("IsInstalled=0")

    $updates = @()
    foreach ($Update in $SearchResult.Updates) {
        $updates += @{
            Title = $Update.Title
            KB = if ($Update.Title -match 'KB(\d+)') { "KB$($Matches[1])" } else { "" }
            Size = [math]::Round($Update.DownloadContents.TotalSize / 1MB, 1)
            Severity = if ($Update.Categories | Where-Object { $_.Name -match 'Critical' }) { "Critical" } else { "Important" }
        }
    }

    $updates | ConvertTo-Json -Depth 3
    '''

    result = subprocess.run(
        ["powershell", "-Command", ps_cmd],
        capture_output=True, text=True, timeout=120
    )

    if result.returncode != 0:
        print(f"錯誤: {result.stderr}")
        return []

    try:
        data = json.loads(result.stdout)
        return data if isinstance(data, list) else [data]
    except json.JSONDecodeError:
        return []


def install_updates_powershell(kb_list=None):
    """
    通過 PowerShell 安裝更新
    kb_list: 指定要安裝的 KB 編號列表，None=全部安裝
    """
    if kb_list:
        kb_filter = ", ".join(f'"{kb}"' for kb in kb_list)
        ps_cmd = f'''
        $criteria = "IsInstalled=0"
        $Session = New-Object -ComObject Microsoft.Update.Session
        $Searcher = $Session.CreateUpdateSearcher()
        $Result = $Searcher.Search($criteria)

        $Installer = $Session.CreateUpdateInstaller()
        $UpdatesToInstall = New-Object -ComObject Microsoft.Update.UpdateColl

        $kbList = @({kb_filter})
        foreach ($Update in $Result.Updates) {
            $kbMatch = $Update.Title -match 'KB(\d+)'
            if ($kbMatch -and $Matches[1] -in $kbList) {{
                $UpdatesToInstall.Add($Update)
            }}
        }}

        $Installer.Updates = $UpdatesToInstall
        $InstallResult = $Installer.Install()
        $InstallResult | ConvertTo-Json
        '''
    else:
        ps_cmd = r'''
        $Session = New-Object -ComObject Microsoft.Update.Session
        $Searcher = $Session.CreateUpdateSearcher()
        $Result = $Searcher.Search("IsInstalled=0 AND IsDownloaded=1")

        $Installer = $Session.CreateUpdateInstaller()
        $Installer.Updates = $Result.Updates
        $InstallResult = $Installer.Install()

        @{
            Installed = $InstallResult.ResultCode
            RebootRequired = $Installer.RebootRequired
            UpdatesCount = $Result.Updates.Count
        } | ConvertTo-Json
        '''

    result = subprocess.run(
        ["powershell", "-Command", ps_cmd],
        capture_output=True, text=True, timeout=600
    )

    print(result.stdout)
    return result.returncode == 0

import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

def remote_patch_scan(
    computer,
    username=None,
    password=None
):
    """
    通過 PowerShell Remoting 遠程掃描補丁
    """
    cred_part = ""
    if username and password:
        cred_part = (
            f'$secpass = ConvertTo-SecureString "{password}" -AsPlainText -Force; '
            f'$cred = New-Object System.Management.Automation.PSCredential("{username}", $secpass); '
        )

    ps_script = f'''
    {cred_part}
    $session = New-PSSession -ComputerName "{computer}" {f"-Credential $cred" if username else ""}
    Invoke-Command -Session $session -ScriptBlock {{
        $Searcher = New-Object -ComObject Microsoft.Update.Session
        $Searcher = $Searcher.CreateUpdateSearcher()
        $Result = $Searcher.Search("IsInstalled=0")

        @{{
            Computer = $env:COMPUTERNAME
            PendingUpdates = $Result.Updates.Count
            CriticalCount = ($Result.Updates | Where-Object {{
                $_.Categories.Name -match 'Critical'
            }}).Count
            LastScanDate = (Get-HotFix | Sort-Object InstalledOn -Descending | Select-Object -First 1).InstalledOn
        }} | ConvertTo-Json
    }}
    Remove-PSSession -Session $session -ErrorAction SilentlyContinue
    '''

    result = subprocess.run(
        ["powershell", "-Command", ps_script],
        capture_output=True, text=True, timeout=120
    )

    try:
        data = json.loads(result.stdout.strip().split('\n')[-1])
        return {"computer": computer, "status": "ok", **data}
    except (json.JSONDecodeError, IndexError):
        return {
            "computer": computer,
            "status": "error",
            "error": result.stderr.strip()
        }


def batch_scan_patches(
    computers,
    max_workers=10,
    username=None,
    password=None
):
    """批量掃描多臺電腦的補丁狀態(tài)"""
    results = []

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(
                remote_patch_scan, pc, username, password
            ): pc for pc in computers
        }

        for future in as_completed(futures):
            result = future.result()
            results.append(result)
            icon = "?" if result["status"] == "ok" else "?"
            if result["status"] == "ok":
                pending = result.get("PendingUpdates", "?")
                critical = result.get("CriticalCount", "?")
                print(f"  {icon} {result['computer']}: "
                      f"{pending} 待安裝, {critical} 關(guān)鍵")
            else:
                print(f"  {icon} {result['computer']}: {result.get('error', '未知錯誤')}")

    return results


def generate_compliance_report(results, output_file="patch_compliance.html"):
    """生成補丁合規(guī)報告"""
    total = len(results)
    compliant = sum(
        1 for r in results
        if r["status"] == "ok" and r.get("PendingUpdates", 1) == 0
    )
    non_compliant = total - compliant

    # 超過 30 天未更新的機器
    stale = []
    for r in results:
        if r["status"] == "ok":
            last_scan = r.get("LastScanDate")
            if last_scan:
                try:
                    scan_date = datetime.strptime(
                        str(last_scan)[:10], "%Y-%m-%d"
                    )
                    if (datetime.now() - scan_date).days > 30:
                        stale.append(r["computer"])
                except (ValueError, TypeError):
                    pass

    compliance_rate = (compliant / total * 100) if total > 0 else 0

    html = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>補丁合規(guī)報告</title>
    <style>
        body {{ font-family: 'Segoe UI', sans-serif; background: #f5f5f5; padding: 20px; }}
        .container {{ max-width: 1200px; margin: 0 auto; }}
        h1 {{ color: #333; }}
        .summary {{ display: grid; grid-template-columns: repeat(3, 1fr); gap: 20px; margin: 20px 0; }}
        .card {{ background: white; border-radius: 8px; padding: 20px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); text-align: center; }}
        .card .number {{ font-size: 36px; font-weight: bold; }}
        .card.good .number {{ color: #27ae60; }}
        .card.bad .number {{ color: #e74c3c; }}
        .card.warn .number {{ color: #f39c12; }}
        table {{ width: 100%; background: white; border-radius: 8px; overflow: hidden; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }}
        th {{ background: #2c3e50; color: white; padding: 12px; text-align: left; }}
        td {{ padding: 10px 12px; border-bottom: 1px solid #eee; }}
        tr:hover {{ background: #f9f9f9; }}
        .status-ok {{ color: #27ae60; font-weight: bold; }}
        .status-warn {{ color: #e74c3c; font-weight: bold; }}
    </style>
</head>
<body>
    <div class="container">
        <h1>Windows 補丁合規(guī)報告</h1>
        <p>生成時間: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>

        <div class="summary">
            <div class="card good">
                <div class="number">{compliance_rate:.0f}%</div>
                <div>合規(guī)率</div>
            </div>
            <div class="card good">
                <div class="number">{compliant}</div>
                <div>已合規(guī)</div>
            </div>
            <div class="card bad">
                <div class="number">{non_compliant}</div>
                <div>待處理</div>
            </div>
        </div>

        <table>
            <tr>
                <th>計算機名</th>
                <th>待安裝補丁</th>
                <th>關(guān)鍵補丁</th>
                <th>最后掃描日期</th>
                <th>狀態(tài)</th>
            </tr>"""

    for r in sorted(results, key=lambda x: x.get("PendingUpdates", 0), reverse=True):
        if r["status"] == "ok":
            pending = r.get("PendingUpdates", 0)
            critical = r.get("CriticalCount", 0)
            last_scan = str(r.get("LastScanDate", "未知"))[:10]
            status_class = "status-warn" if pending > 0 else "status-ok"
            status_text = "待處理" if pending > 0 else "合規(guī)"
            html += f"""
            <tr>
                <td>{r['computer']}</td>
                <td>{pending}</td>
                <td>{critical}</td>
                <td>{last_scan}</td>
                <td class="{status_class}">{status_text}</td>
            </tr>"""

    html += """
        </table>
    </div>
</body>
</html>"""

    with open(output_file, "w", encoding="utf-8") as f:
        f.write(html)

    print(f"\n合規(guī)報告已生成: {output_file}")
    print(f"合規(guī)率: {compliance_rate:.1f}% ({compliant}/{total})")
    return output_file


# 使用示例
# computers = [f"PC-{i:03d}" for i in range(1, 101)]
# results = batch_scan_patches(computers)
# generate_compliance_report(results)

def uninstall_update(kb_number):
    """
    卸載指定的 KB 補丁
    kb_number: 例如 "KB5005565"
    """
    # 方法1：使用 wusa 命令
    cmd = f"wusa /uninstall /kb:{kb_number.replace('KB', '')} /quiet /norestart"

    result = subprocess.run(
        cmd, shell=True, capture_output=True, text=True, timeout=300
    )

    if result.returncode == 0:
        print(f"補丁 {kb_number} 已卸載")
        print("建議重啟電腦以完成卸載")
        return True
    else:
        print(f"卸載失敗: {result.stderr}")
        return False


def list_installed_hotfixes(search_term=""):
    """列出已安裝的補丁"""
    ps_cmd = f'Get-HotFix | Where-Object {{ $_.HotFixID -like "*{search_term}*" }} | Sort-Object InstalledOn -Descending | Select-Object -First 20 | Format-Table -AutoSize'

    result = subprocess.run(
        ["powershell", "-Command", ps_cmd],
        capture_output=True, text=True, timeout=30
    )

    print(result.stdout)
    return result.stdout


def check_vulnerability_status():
    """
    檢查系統(tǒng)的已知漏洞狀態(tài)
    通過對比已安裝補丁和已知漏洞 KB 列表
    """
    # 常見的關(guān)鍵安全補丁列表（示例）
    critical_kbs = [
        ("KB5005565", "PrintNightmare 打印機漏洞"),
        ("KB5012170", "CVE-2022-26925 LSA 泄露"),
        ("KB5009543", "Active Directory 安全繞過"),
        ("KB5011048", "遠程代碼執(zhí)行漏洞"),
    ]

    print("檢查關(guān)鍵漏洞補丁狀態(tài)：\n")

    # 獲取已安裝補丁列表
    result = subprocess.run(
        ["powershell", "-Command", "Get-HotFix | Select-Object HotFixID"],
        capture_output=True, text=True, timeout=30
    )
    installed = set(result.stdout.splitlines())

    for kb, desc in critical_kbs:
        status = "? 已安裝" if kb in installed else "? 未安裝"
        color = "安全" if kb in installed else "?? 有風(fēng)險"
        print(f"  {status} | {kb} | {desc} | {color}")

    return critical_kbs

def staged_deployment(computers, batch_size=10, delay_hours=24):
    """
    分批部署補丁
    先在第一批驗證，確認無問題后再推后續(xù)批次
    """
    batches = [
        computers[i:i + batch_size]
        for i in range(0, len(computers), batch_size)
    ]

    print(f"共 {len(computers)} 臺電腦，分為 {len(batches)} 批")
    print(f"每批 {batch_size} 臺，間隔 {delay_hours} 小時\n")

    for i, batch in enumerate(batches, 1):
        print(f"=== 第 {i} 批 ({len(batch)} 臺) ===")
        print(f"電腦: {', '.join(batch)}")

        # 第一批需要人工確認
        if i == 1:
            print("?? 這是第一批，建議先在此批次驗證補丁兼容性")
            confirm = input("確認安裝？(y/n): ").strip().lower()
            if confirm != 'y':
                print("跳過此批次")
                continue

        # 執(zhí)行安裝（此處調(diào)用實際的安裝函數(shù)）
        print(f"正在安裝...\n")

        # 實際場景中，這里調(diào)用 remote_patch_install
        # 如果不是最后一批，等待確認
        if i < len(batches):
            print(f"第 {i} 批完成。等待 {delay_hours} 小時后繼續(xù)下一批")
            print("期間請驗證業(yè)務(wù)系統(tǒng)正常運行")

PATCH_POLICY = {
    "Critical": {
        "deploy_window": "72小時內(nèi)",     # 關(guān)鍵補丁 3 天內(nèi)部署
        "reboot_required": True,
        "approval": "自動",
    },
    "Important": {
        "deploy_window": "14天內(nèi)",       # 重要補丁 2 周內(nèi)部署
        "reboot_required": True,
        "approval": "測試后自動",
    },
    "Moderate": {
        "deploy_window": "30天內(nèi)",       # 中等補丁 1 個月內(nèi)
        "reboot_required": False,
        "approval": "人工審核",
    },
    "Low": {
        "deploy_window": "隨下次更新周期",  # 低風(fēng)險隨緣
        "reboot_required": False,
        "approval": "人工審核",
    },
}