Alibabacloud Sas Multiaccount Manage

#!/usr/bin/env python3
"""accounts.py — 多账号管理工具

用法:
  uv run accounts.py refresh
  uv run accounts.py search <DisplayName>
  uv run accounts.py enable <AccountId>
  uv run accounts.py disable <AccountId>
  uv run accounts.py list
"""

import argparse
import json
import subprocess
import sys
from pathlib import Path

ACCOUNTS_FILE = Path("accounts.json")
ALIYUN_USER_AGENT_HEADER = "User-Agent=AlibabaCloud-Agent-Skills/alibabacloud-sas-multiaccount-manage"
CLI_CONNECT_TIMEOUT_SECONDS = 10
CLI_READ_TIMEOUT_SECONDS = 60


def _aliyun_cmd(*args):
    """构建统一的 aliyun CLI 参数（含 User-Agent 与超时配置）。"""
    return [
        "aliyun",
        "--header",
        ALIYUN_USER_AGENT_HEADER,
        "--connect-timeout",
        str(CLI_CONNECT_TIMEOUT_SECONDS),
        "--read-timeout",
        str(CLI_READ_TIMEOUT_SECONDS),
        *args,
    ]


def load_accounts():
    if not ACCOUNTS_FILE.exists():
        print("错误: accounts.json 不存在，请先执行 refresh", file=sys.stderr)
        sys.exit(1)
    return json.loads(ACCOUNTS_FILE.read_text(encoding="utf-8"))


def save_accounts(accounts):
    ACCOUNTS_FILE.write_text(
        json.dumps(accounts, indent=2, ensure_ascii=False), encoding="utf-8"
    )


def cmd_refresh(args):
    """调用 aliyun sas list-accounts-in-resource-directory，写入 accounts.json"""
    region_id = getattr(args, "region_id", "cn-shanghai")

    # 获取当前凭证的主账号（自身也应包含在可操作范围内）
    identity_result = subprocess.run(
        _aliyun_cmd("sts", "get-caller-identity"),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )
    if identity_result.returncode != 0:
        print(f"get-caller-identity 调用失败:\n{identity_result.stderr}", file=sys.stderr)
        sys.exit(1)
    caller_account_id = str(json.loads(identity_result.stdout)["AccountId"])

    result = subprocess.run(
        _aliyun_cmd("sas", "--region", region_id, "list-accounts-in-resource-directory"),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )
    if result.returncode != 0:
        print(f"aliyun CLI 调用失败:\n{result.stderr}", file=sys.stderr)
        sys.exit(1)

    data = json.loads(result.stdout)
    accounts = data.get("Accounts", [])

    # 调用 describe-monitor-accounts，只保留已纳入监控的账号（并将自身主账号也包含进来）
    monitor_result = subprocess.run(
        _aliyun_cmd("sas", "--region", region_id, "describe-monitor-accounts"),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )
    if monitor_result.returncode != 0:
        print(f"describe-monitor-accounts 调用失败:\n{monitor_result.stderr}", file=sys.stderr)
        sys.exit(1)

    monitor_data = json.loads(monitor_result.stdout)
    monitored_ids = set(str(i) for i in monitor_data.get("AccountIds", []))
    monitored_ids.add(caller_account_id)  # 自身主账号始终属于可操作范围

    accounts = [a for a in accounts if str(a["AccountId"]) in monitored_ids]

    # 保留已有的 enable 状态，新账号默认 enable=true
    existing = {}
    if ACCOUNTS_FILE.exists():
        for a in json.loads(ACCOUNTS_FILE.read_text(encoding="utf-8")):
            existing[a["AccountId"]] = a.get("enable", True)

    for account in accounts:
        account["enable"] = existing.get(account["AccountId"], True)

    save_accounts(accounts)
    print(f"已刷新 {len(accounts)} 个账号，写入 {ACCOUNTS_FILE}")


def cmd_search(args):
    """按 DisplayName 模糊搜索，输出 AccountId"""
    keyword = args.keyword.lower()
    accounts = load_accounts()
    results = [
        a for a in accounts if keyword in a.get("DisplayName", "").lower()
    ]
    if not results:
        print(f"未找到匹配 '{args.keyword}' 的账号")
        return
    for a in results:
        status = "启用" if a.get("enable", True) else "禁用"
        print(f"{a['AccountId']}\t{a.get('DisplayName', '')}\t[{status}]")


def _set_enable(account_id, value):
    accounts = load_accounts()
    found = False
    for a in accounts:
        if a["AccountId"] == account_id:
            a["enable"] = value
            found = True
            break
    if not found:
        print(f"错误: 未找到账号 {account_id}", file=sys.stderr)
        sys.exit(1)
    save_accounts(accounts)
    action = "启用" if value else "禁用"
    print(f"账号 {account_id} 已{action}")


def cmd_enable(args):
    _set_enable(args.account_id, True)


def cmd_disable(args):
    _set_enable(args.account_id, False)


def cmd_list(_args):
    """列出所有账号"""
    accounts = load_accounts()
    for a in accounts:
        status = "启用" if a.get("enable", True) else "禁用"
        print(f"{a['AccountId']}\t{a.get('DisplayName', ''):<20}\t[{status}]")


def get_enabled_accounts():
    """供其他模块调用：返回所有 enable=True 的账号列表"""
    return [a for a in load_accounts() if a.get("enable", True)]


def get_caller_account_id():
    """获取当前凭证的主账号 ID（通过 aliyun sts get-caller-identity）。"""
    result = subprocess.run(
        _aliyun_cmd("sts", "get-caller-identity"),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )
    if result.returncode != 0:
        print(f"get-caller-identity 调用失败:\n{result.stderr}", file=sys.stderr)
        sys.exit(1)
    return str(json.loads(result.stdout)["AccountId"])


def main():
    parser = argparse.ArgumentParser(
        description="阿里云云安全中心多账号管理工具"
    )
    sub = parser.add_subparsers(dest="command", metavar="command")
    sub.required = True

    p_refresh = sub.add_parser("refresh", help="刷新账号列表（从资源目录拉取）")
    p_refresh.add_argument(
        "--region-id",
        dest="region_id",
        choices=["cn-shanghai", "ap-southeast-1"],
        default="cn-shanghai",
        help="SAS API 地域：cn-shanghai（中国大陆，默认）/ ap-southeast-1（非中国大陆）",
    )

    p_search = sub.add_parser("search", help="按 DisplayName 搜索账号")
    p_search.add_argument("keyword", help="搜索关键字")

    p_enable = sub.add_parser("enable", help="启用指定账号")
    p_enable.add_argument("account_id", help="账号 ID")

    p_disable = sub.add_parser("disable", help="禁用指定账号")
    p_disable.add_argument("account_id", help="账号 ID")

    sub.add_parser("list", help="列出所有账号及状态")

    args = parser.parse_args()
    dispatch = {
        "refresh": cmd_refresh,
        "search": cmd_search,
        "enable": cmd_enable,
        "disable": cmd_disable,
        "list": cmd_list,
    }
    dispatch[args.command](args)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""baseline.py — 云安全中心基线/系统基线批量导出工具

用法:
  # 导出所有启用账号的云平台配置检查结果
  uv run baseline.py export-cspm

  # 仅导出指定账号
  uv run baseline.py export-cspm --account-id 1234567890

  # 导出系统基线风险列表
  uv run baseline.py export-system-warning
  uv run baseline.py export-system-warning --account-id 1234567890
"""

import argparse
import asyncio
import json
import shutil
import sys
import urllib.request
import zipfile
from datetime import date
from pathlib import Path

# 将 scripts 目录加入路径以便导入同级模块
sys.path.insert(0, str(Path(__file__).parent))
from accounts import get_caller_account_id, get_enabled_accounts  # noqa: E402
from merge import merge_excel  # noqa: E402

TODAY = date.today().strftime("%Y%m%d")
QPS_LIMIT = 5  # API 并发上限
ALIYUN_USER_AGENT_HEADER = "User-Agent=AlibabaCloud-Agent-Skills/alibabacloud-sas-multiaccount-manage"
CLI_CONNECT_TIMEOUT_SECONDS = 10
CLI_READ_TIMEOUT_SECONDS = 60
DOWNLOAD_TIMEOUT_SECONDS = 60

# 全局信号量，在事件循环内初始化

# 当前凭证的主账号 ID，在 do_export() 中初始化
_caller_account_id = ""

# SAS API 地域，在 do_export() 中初始化
_region_id = "cn-shanghai"

# 可跳过的 API 错误码（账号无权限/免费版限制等），静默忽略
_SKIP_ERROR_CODES = {"FreeVersionNotPermit", "NoPermission", "Forbidden"}


class AccountSkippedError(Exception):
    """账号因权限不足等原因被跳过，不中断整体流程。"""
    def __init__(self, account_id, reason):
        self.account_id = account_id
        self.reason = reason
        super().__init__(f"账号 {account_id} 跳过: {reason}")

# ────────────────────────────── 异步 API 封装 ──────────────────────────────


async def _run_aliyun_async(args, account_id=""):
    """异步运行 aliyun CLI，通过信号量将并发 API 调用限制在 QPS ≤ 5。

    对可跳过的错误码（如 FreeVersionNotPermit）抛出 AccountSkippedError。
    """
    async with _api_sem:
        proc = await asyncio.create_subprocess_exec(
            "aliyun",
            "--header",
            ALIYUN_USER_AGENT_HEADER,
            "--connect-timeout",
            str(CLI_CONNECT_TIMEOUT_SECONDS),
            "--read-timeout",
            str(CLI_READ_TIMEOUT_SECONDS),
            "sas",
            "--region",
            _region_id,
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await proc.communicate()

    err_text = stderr.decode()
    if proc.returncode != 0:
        # 检查是否属于可跳过的错误码（静默忽略）
        for code in _SKIP_ERROR_CODES:
            if code in err_text:
                raise AccountSkippedError(account_id, code)
        print(f"aliyun CLI 调用失败:\n{err_text}", file=sys.stderr)
        raise RuntimeError("aliyun CLI error")
    # returncode==0 但 stderr 中仍含可跳过错误码时也忽略
    for code in _SKIP_ERROR_CODES:
        if code in err_text:
            raise AccountSkippedError(account_id, code)
    try:
        return json.loads(stdout.decode())
    except json.JSONDecodeError:
        print(f"响应解析失败:\n{stdout.decode()}", file=sys.stderr)
        raise


async def start_export_async(export_type, account_id, params=None):
    """发起导出任务，返回 export_id。"""
    cli_args = [
        "export-record",
        "--lang", "zh",
        "--export-type", export_type,
    ]
    if account_id != _caller_account_id:
        cli_args += ["--resource-directory-account-id", str(account_id)]
    if params:
        cli_args += ["--params", params]

    data = await _run_aliyun_async(cli_args, account_id)
    export_id = str(data["Id"])
    print(f"  [提交] 账号 {account_id}，export_id={data['Id']}")
    return export_id


async def wait_for_export_async(export_id, account_id, poll_interval=5):
    """轮询导出状态，成功后返回下载链接。"""
    while True:
        cli_args = [
            "describe-export-info",
            "--export-id", export_id,
        ]
        if account_id != _caller_account_id:
            cli_args += ["--resource-directory-account-id", str(account_id)]
        data = await _run_aliyun_async(cli_args, account_id)
        status = data.get("ExportStatus", "")

        if status == "success":
            return data["Link"]
        if status in ("failed", "error"):
            raise RuntimeError(f"账号 {account_id} 导出失败，状态={status}")
        await asyncio.sleep(poll_interval)


async def download_and_extract_async(link, account_id, prefix):
    """异步下载 zip，解压，重命名为 {prefix}-{account_id}-{date}.xlsx。"""
    zip_path = Path(f"{prefix}-{account_id}-{TODAY}.zip")
    print(f"  [下载] 账号 {account_id}...")

    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, _download_with_timeout, link, zip_path)

    with zipfile.ZipFile(zip_path) as zf:
        names = zf.namelist()
        zf.extractall(path=zip_path.parent)
        extracted = zip_path.parent / names[0]

    output_path = Path(f"{prefix}-{account_id}-{TODAY}.xlsx")
    if extracted != output_path:
        extracted.rename(output_path)
    if zip_path.exists():
        zip_path.unlink()

    return str(output_path)


def _download_with_timeout(link, output_path):
    """下载文件并设置超时，避免网络异常时无限阻塞。"""
    with urllib.request.urlopen(link, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response:
        with Path(output_path).open("wb") as fp:
            shutil.copyfileobj(response, fp)


async def _wait_and_download(export_id, account_id, prefix):
    """等待导出完成后下载，返回 merge 所需的 input dict。"""
    link = await wait_for_export_async(export_id, account_id)
    xlsx_path = await download_and_extract_async(link, account_id, prefix)
    return {"filename": xlsx_path, "account_id": account_id}


# ────────────────────────────── 异步导出流程 ──────────────────────────────


async def do_export_async(export_type, prefix, merged_name, account_ids, params=None):
    """并发导出流程：

    阶段 1 — 并发提交所有账号的导出任务（QPS ≤ 5）
    阶段 2 — 并发轮询 + 下载解压（QPS ≤ 5）
    阶段 3 — 合并所有 xlsx
    """
    global _api_sem
    _api_sem = asyncio.Semaphore(QPS_LIMIT)

    if not account_ids:
        print("错误: 没有可用账号", file=sys.stderr)
        sys.exit(1)

    print(f"共 {len(account_ids)} 个账号待导出（耗时操作，请耐心等待）")

    # 阶段 1：并发提交所有导出任务
    submit_results = await asyncio.gather(
        *[start_export_async(export_type, aid, params) for aid in account_ids],
        return_exceptions=True,
    )

    pending = []  # (export_id, account_id)
    skipped = []
    for aid, res in zip(account_ids, submit_results):
        if isinstance(res, AccountSkippedError):
            skipped.append(aid)
        elif isinstance(res, BaseException):
            print(f"  [失败] 账号 {aid}: {res}", file=sys.stderr)
        else:
            pending.append((res, aid))

    if not pending:
        print(f"所有账号均被跳过（{len(skipped)} 个），无可合并数据")
        return

    # 阶段 2：并发等待 + 下载解压
    download_results = await asyncio.gather(
        *[_wait_and_download(eid, aid, prefix) for eid, aid in pending],
        return_exceptions=True,
    )

    inputs = []
    failed = []
    for (_, aid), res in zip(pending, download_results):
        if isinstance(res, AccountSkippedError):
            skipped.append(aid)
        elif isinstance(res, BaseException):
            print(f"  [失败] 账号 {aid}: {res}", file=sys.stderr)
            failed.append(aid)
        else:
            print(f"  [成功] 账号 {aid}")
            inputs.append(res)

    if skipped:
        print(f"跳过 {len(skipped)} 个账号: {', '.join(skipped)}")
    if failed:
        print(f"失败 {len(failed)} 个账号: {', '.join(failed)}", file=sys.stderr)

    if not inputs:
        print("没有成功下载的文件，跳过合并")
        return

    # 阶段 3：合并
    merge_excel(merged_name, inputs)

    # 阶段 4：删除临时 xlsx 文件
    for item in inputs:
        tmp = Path(item["filename"])
        if tmp.exists():
            tmp.unlink()
    print(f"已生成: {merged_name}（共 {len(inputs)} 个账号，已清理临时文件）")


def do_export(export_type, prefix, merged_name, account_ids, params=None, region_id="cn-shanghai"):
    """同步入口，内部通过 asyncio.run 驱动异步流程。"""
    global _caller_account_id, _region_id
    _caller_account_id = get_caller_account_id()
    _region_id = region_id
    loop = asyncio.get_event_loop()
    loop.run_until_complete(
        do_export_async(export_type, prefix, merged_name, account_ids, params)
    )


# ────────────────────────────── CLI 子命令 ──────────────────────────────


def cmd_export_cspm(args):
    """导出云平台配置检查结果（baselineCspm）。"""
    if args.account_id:
        account_ids = [args.account_id]
    else:
        accounts = get_enabled_accounts()
        account_ids = [a["AccountId"] for a in accounts]

    do_export(
        export_type="baselineCspm",
        prefix="baseline-cspm",
        merged_name=f"baseline-cspm-merged-{TODAY}.xlsx",
        account_ids=account_ids,
        region_id=args.region_id,
    )


def cmd_export_system_warning(args):
    """导出系统基线风险列表（exportHcWarning）。"""
    if args.account_id:
        account_ids = [args.account_id]
    else:
        accounts = get_enabled_accounts()
        account_ids = [a["AccountId"] for a in accounts]

    params = json.dumps(
        {
            "CheckLevel": "high,medium,low",
            "CheckWarningStatusList": [1, 3, 6, 8],
            "Source": "default",
        },
        ensure_ascii=False,
    )

    do_export(
        export_type="exportHcWarning",
        prefix="system-warning",
        merged_name=f"system-warning-merged-{TODAY}.xlsx",
        account_ids=account_ids,
        params=params,
        region_id=args.region_id,
    )


# ────────────────────────────── 入口 ──────────────────────────────


def main():
    parser = argparse.ArgumentParser(description="云安全中心基线批量导出工具")
    sub = parser.add_subparsers(dest="command", metavar="command")
    sub.required = True

    # export-cspm
    p_cspm = sub.add_parser(
        "export-cspm",
        help="导出云平台配置检查结果（baselineCspm）",
    )
    p_cspm.add_argument(
        "--account-id",
        metavar="ACCOUNT_ID",
        help="指定单个账号 ID（默认导出所有启用账号）",
    )
    p_cspm.add_argument(
        "--region-id",
        dest="region_id",
        choices=["cn-shanghai", "ap-southeast-1"],
        default="cn-shanghai",
        help="SAS API 地域：cn-shanghai（中国大陆，默认）/ ap-southeast-1（非中国大陆）",
    )

    # export-system-warning
    p_warn = sub.add_parser(
        "export-system-warning",
        help="导出系统基线风险列表（exportHcWarning）",
    )
    p_warn.add_argument(
        "--account-id",
        metavar="ACCOUNT_ID",
        help="指定单个账号 ID（默认导出所有启用账号）",
    )
    p_warn.add_argument(
        "--region-id",
        dest="region_id",
        choices=["cn-shanghai", "ap-southeast-1"],
        default="cn-shanghai",
        help="SAS API 地域：cn-shanghai（中国大陆，默认）/ ap-southeast-1（非中国大陆）",
    )

    args = parser.parse_args()
    dispatch = {
        "export-cspm": cmd_export_cspm,
        "export-system-warning": cmd_export_system_warning,
    }
    dispatch[args.command](args)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""merge.py — 公共 Excel 表格合并工具

用法（作为模块导入）:
    from merge import merge_excel
    merge_excel(
        export_name="output.xlsx",
        inputs=[
            {"filename": "a.xlsx", "account_id": "123456"},
            {"filename": "b.xlsx", "account_id": "789012"},
        ],
    )

用法（命令行）:
    uv run merge.py --output merged.xlsx --input a.xlsx:123 b.xlsx:456
"""

import argparse
import sys
import warnings
from pathlib import Path

import openpyxl

warnings.filterwarnings(
    "ignore",
    message="Workbook contains no default style",
    category=UserWarning,
)


def merge_excel(export_name, inputs):
    """将多个 Excel 文件合并为一个，并在末尾追加「资源管理账号」列。

    Args:
        export_name: 输出文件名（.xlsx）
        inputs: 列表，每项包含 filename(str) 和 account_id(str)

    Returns:
        输出文件的路径字符串
    """
    if not inputs:
        raise ValueError("inputs 不能为空")

    all_rows = []
    header = None
    valid_count = 0

    for item in inputs:
        filename = item["filename"]
        account_id = str(item["account_id"])
        if not Path(filename).exists():
            print("警告: 文件 {} 不存在，跳过".format(filename), file=sys.stderr)
            continue

        wb = openpyxl.load_workbook(filename, data_only=True)
        ws = wb["Sheet0"] if "Sheet0" in wb.sheetnames else wb.active
        rows = list(ws.values)
        wb.close()

        if not rows:
            continue

        if header is None:
            header = list(rows[0]) + ["资源管理账号"]
            all_rows.append(header)

        for row in rows[1:]:
            all_rows.append(list(row) + [account_id])

        valid_count += 1

    if not all_rows or valid_count == 0:
        raise RuntimeError("没有可合并的有效文件")

    out_wb = openpyxl.Workbook()
    out_ws = out_wb.active
    out_ws.title = "Sheet0"
    for row in all_rows:
        out_ws.append(row)
    out_wb.save(export_name)

    data_rows = len(all_rows) - 1
    print("合并完成: {}（共 {} 行，{} 个账号）".format(export_name, data_rows, valid_count))
    return export_name


def main():
    parser = argparse.ArgumentParser(description="多账号 Excel 表格合并工具")
    parser.add_argument("--output", "-o", required=True, help="输出文件名")
    parser.add_argument(
        "--input",
        "-i",
        nargs="+",
        required=True,
        metavar="FILE:ACCOUNT_ID",
        help="输入文件，格式: 文件路径:账号ID，可指定多个",
    )
    args = parser.parse_args()

    inputs = []
    for item in args.input:
        parts = item.rsplit(":", 1)
        if len(parts) != 2:
            print(
                "错误: 输入格式不正确 '{}'，应为 '文件路径:账号ID'".format(item),
                file=sys.stderr,
            )
            sys.exit(1)
        inputs.append({"filename": parts[0], "account_id": parts[1]})

    merge_excel(args.output, inputs)


if __name__ == "__main__":
    main()

[project]
name = "scripts"
version = "0.1.0"
description = "阿里云云安全中心多账号管理工具"
requires-python = ">=3.6"
dependencies = [
    "openpyxl>=3.1.5",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
only-include = [
    "accounts.py",
    "merge.py",
    "baseline.py",
    "export_baseline.py",
    "vuln.py",
]

[project.scripts]
accounts = "accounts:main"
merge = "merge:main"
baseline = "baseline:main"
vuln = "vuln:main"

et_xmlfile==2.0.0
numpy==2.4.4
openpyxl==3.1.5
python-dateutil==2.9.0.post0
six==1.17.0

#!/usr/bin/env python3
"""vuln.py — 云安全中心漏洞批量导出工具

用法:
  uv run vuln.py export-cve              # Linux 软件漏洞
  uv run vuln.py export-sys              # Windows 系统漏洞
  uv run vuln.py export-app              # 应用漏洞（含 SCA）
  uv run vuln.py export-emg              # 应急漏洞
  uv run vuln.py export-all              # 依次导出全部四种类型

  # 仅导出指定账号
  uv run vuln.py export-cve --account-id 1234567890
"""

import argparse
import asyncio
import json
import shutil
import sys
import urllib.request
import zipfile
from datetime import date
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from accounts import get_caller_account_id, get_enabled_accounts  # noqa: E402
from merge import merge_excel  # noqa: E402

TODAY = date.today().strftime("%Y%m%d")
QPS_LIMIT = 5
ALIYUN_USER_AGENT_HEADER = "User-Agent=AlibabaCloud-Agent-Skills/alibabacloud-sas-multiaccount-manage"
CLI_CONNECT_TIMEOUT_SECONDS = 10
CLI_READ_TIMEOUT_SECONDS = 60
DOWNLOAD_TIMEOUT_SECONDS = 60

_api_sem = None  # 在事件循环内初始化

# 当前凭证的主账号 ID，在 do_export() 中初始化
_caller_account_id = ""

# SAS API 地域，在 do_export() 中初始化
_region_id = "cn-shanghai"

_SKIP_ERROR_CODES = {"FreeVersionNotPermit", "NoPermission", "Forbidden"}

# ────────────────────────────── 导出类型配置 ──────────────────────────────

EXPORT_CONFIGS = {
    "cve": {
        "cli_args": [
            "--lang", "zh",
            "--type", "cve",
            "--necessity", "asap,later,nntf",
            "--dealed", "n",
        ],
        "prefix": "vul-cve",
        "desc": "Linux 软件漏洞",
    },
    "sys": {
        "cli_args": [
            "--lang", "zh",
            "--type", "sys",
            "--necessity", "asap,later,nntf",
            "--dealed", "n",
        ],
        "prefix": "vul-sys",
        "desc": "Windows 系统漏洞",
    },
    "app": {
        "cli_args": [
            "--lang", "zh",
            "--type", "app",
            "--necessity", "asap,later,nntf",
            "--attach-types", "sca",
            "--asset-type", "ECS,CONTAINER",
            "--dealed", "n",
        ],
        "prefix": "vul-app",
        "desc": "应用漏洞",
    },
    "emg": {
        "cli_args": [
            "--lang", "zh",
            "--type", "emg",
            "--risk-status", "y",
            "--dealed", "n",
        ],
        "prefix": "vul-emg",
        "desc": "应急漏洞",
    },
}


# ────────────────────────────── 异常 ──────────────────────────────


class AccountSkippedError(Exception):
    """账号因权限不足等原因被跳过，不中断整体流程。"""

    def __init__(self, account_id, reason):
        self.account_id = account_id
        self.reason = reason
        super().__init__(f"账号 {account_id} 跳过: {reason}")


# ────────────────────────────── 异步 API 封装 ──────────────────────────────


async def _run_aliyun_async(args, account_id=""):
    """异步运行 aliyun CLI，通过信号量将并发 API 调用限制在 QPS ≤ 5。"""
    async with _api_sem:
        proc = await asyncio.create_subprocess_exec(
            "aliyun",
            "--header",
            ALIYUN_USER_AGENT_HEADER,
            "--connect-timeout",
            str(CLI_CONNECT_TIMEOUT_SECONDS),
            "--read-timeout",
            str(CLI_READ_TIMEOUT_SECONDS),
            "sas",
            "--region",
            _region_id,
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await proc.communicate()

    err_text = stderr.decode()
    if proc.returncode != 0:
        for code in _SKIP_ERROR_CODES:
            if code in err_text:
                raise AccountSkippedError(account_id, code)
        print(f"aliyun CLI 调用失败:\n{err_text}", file=sys.stderr)
        raise RuntimeError("aliyun CLI error")
    # returncode==0 但 stderr 中仍含可跳过错误码时也忽略
    for code in _SKIP_ERROR_CODES:
        if code in err_text:
            raise AccountSkippedError(account_id, code)
    try:
        return json.loads(stdout.decode())
    except json.JSONDecodeError:
        print(f"响应解析失败:\n{stdout.decode()}", file=sys.stderr)
        raise


async def start_export_async(vul_type, account_id):
    """发起漏洞导出任务，返回 export_id。"""
    cli_args = ["export-vul", "--force"] + EXPORT_CONFIGS[vul_type]["cli_args"]
    if account_id != _caller_account_id:
        cli_args += ["--resource-directory-account-id", str(account_id)]

    data = await _run_aliyun_async(cli_args, account_id)
    print(f"  [提交] 账号 {account_id}，export_id={data['Id']}")
    return str(data["Id"])


async def wait_for_export_async(export_id, account_id, poll_interval=5):
    """轮询 describe-vul-export-info，成功后返回下载链接。"""
    while True:
        cli_args = [
            "describe-vul-export-info",
            "--force",
            "--export-id", export_id,
        ]
        if account_id != _caller_account_id:
            cli_args += ["--resource-directory-account-id", str(account_id)]
        data = await _run_aliyun_async(cli_args, account_id)
        status = data.get("ExportStatus", "")
        if status == "success":
            return data["Link"]
        if status in ("failed", "error"):
            raise RuntimeError(f"账号 {account_id} 导出失败，状态={status}")
        await asyncio.sleep(poll_interval)


async def download_and_extract_async(link, account_id, prefix):
    """下载 zip，解压，重命名为 {prefix}-{account_id}-{date}.xlsx。"""
    zip_path = Path(f"{prefix}-{account_id}-{TODAY}.zip")
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, _download_with_timeout, link, zip_path)

    with zipfile.ZipFile(zip_path) as zf:
        names = zf.namelist()
        zf.extractall(path=zip_path.parent)
        extracted = zip_path.parent / names[0]

    output_path = Path(f"{prefix}-{account_id}-{TODAY}.xlsx")
    if extracted != output_path:
        extracted.rename(output_path)
    if zip_path.exists():
        zip_path.unlink()
    return str(output_path)


def _download_with_timeout(link, output_path):
    """下载文件并设置超时，避免网络异常时无限阻塞。"""
    with urllib.request.urlopen(link, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response:
        with Path(output_path).open("wb") as fp:
            shutil.copyfileobj(response, fp)


async def _wait_and_download(export_id, account_id, prefix):
    link = await wait_for_export_async(export_id, account_id)
    xlsx_path = await download_and_extract_async(link, account_id, prefix)
    return {"filename": xlsx_path, "account_id": account_id}


# ────────────────────────────── 导出流程 ──────────────────────────────


async def do_export_async(vul_type, account_ids):
    """单类型并发导出：提交 → 等待 → 下载 → 合并 → 清理临时文件。"""
    global _api_sem
    _api_sem = asyncio.Semaphore(QPS_LIMIT)

    config = EXPORT_CONFIGS[vul_type]
    prefix = config["prefix"]
    merged_name = f"{prefix}-merged-{TODAY}.xlsx"

    print(f"导出【{config['desc']}】共 {len(account_ids)} 个账号（耗时操作，请耐心等待）")

    # 阶段 1：并发提交
    submit_results = await asyncio.gather(
        *[start_export_async(vul_type, aid) for aid in account_ids],
        return_exceptions=True,
    )

    pending = []
    skipped = []
    for aid, res in zip(account_ids, submit_results):
        if isinstance(res, AccountSkippedError):
            skipped.append(aid)
        elif isinstance(res, BaseException):
            print(f"  [失败] 账号 {aid}: {res}", file=sys.stderr)
        else:
            pending.append((res, aid))

    if not pending:
        print(f"所有账号均被跳过（{len(skipped)} 个），跳过【{config['desc']}】")
        return

    # 阶段 2：并发等待 + 下载
    download_results = await asyncio.gather(
        *[_wait_and_download(eid, aid, prefix) for eid, aid in pending],
        return_exceptions=True,
    )

    inputs = []
    failed = []
    for (_, aid), res in zip(pending, download_results):
        if isinstance(res, AccountSkippedError):
            skipped.append(aid)
        elif isinstance(res, BaseException):
            print(f"  [失败] 账号 {aid}: {res}", file=sys.stderr)
            failed.append(aid)
        else:
            print(f"  [成功] 账号 {aid}")
            inputs.append(res)

    if skipped:
        print(f"跳过 {len(skipped)} 个账号: {', '.join(skipped)}")
    if failed:
        print(f"失败 {len(failed)} 个账号: {', '.join(failed)}", file=sys.stderr)

    if not inputs:
        print(f"没有成功下载的文件，跳过【{config['desc']}】合并")
        return

    # 阶段 3：合并
    merge_excel(merged_name, inputs)

    # 阶段 4：清理临时文件
    for item in inputs:
        p = Path(item["filename"])
        if p.exists():
            p.unlink()
    print(f"已生成: {merged_name}（共 {len(inputs)} 个账号，已清理临时文件）")


def do_export(vul_type, account_ids, region_id="cn-shanghai"):
    global _caller_account_id, _region_id
    _caller_account_id = get_caller_account_id()
    _region_id = region_id
    loop = asyncio.get_event_loop()
    loop.run_until_complete(do_export_async(vul_type, account_ids))


# ────────────────────────────── CLI ──────────────────────────────


def _get_account_ids(args):
    if getattr(args, "account_id", None):
        return [args.account_id]
    accounts = get_enabled_accounts()
    if not accounts:
        print("错误: 没有可用账号，请先执行 accounts.py refresh", file=sys.stderr)
        sys.exit(1)
    return [a["AccountId"] for a in accounts]


def cmd_export_cve(args):
    do_export("cve", _get_account_ids(args), args.region_id)


def cmd_export_sys(args):
    do_export("sys", _get_account_ids(args), args.region_id)


def cmd_export_app(args):
    do_export("app", _get_account_ids(args), args.region_id)


def cmd_export_emg(args):
    do_export("emg", _get_account_ids(args), args.region_id)


def main():
    parser = argparse.ArgumentParser(description="云安全中心漏洞批量导出工具")
    sub = parser.add_subparsers(dest="command", metavar="command")
    sub.required = True

    def _add_sub(name, help_text):
        p = sub.add_parser(name, help=help_text)
        p.add_argument(
            "--account-id",
            metavar="ACCOUNT_ID",
            help="指定单个账号 ID（默认导出所有启用账号）",
        )
        p.add_argument(
            "--region-id",
            dest="region_id",
            choices=["cn-shanghai", "ap-southeast-1"],
            default="cn-shanghai",
            help="SAS API 地域：cn-shanghai（中国大陆，默认）/ ap-southeast-1（非中国大陆）",
        )
        return p

    _add_sub("export-cve", "导出 Linux 软件漏洞（cve）")
    _add_sub("export-sys", "导出 Windows 系统漏洞（sys）")
    _add_sub("export-app", "导出应用漏洞（app，含 SCA）")
    _add_sub("export-emg", "导出应急漏洞（emg）")

    args = parser.parse_args()
    dispatch = {
        "export-cve": cmd_export_cve,
        "export-sys": cmd_export_sys,
        "export-app": cmd_export_app,
        "export-emg": cmd_export_emg,
    }
    dispatch[args.command](args)


if __name__ == "__main__":
    main()