xiaoxia-saas/scripts/ci/acr_cleanup.py

#!/usr/bin/env python3
"""
ACR 镜像清理脚本（增强版）

清理策略：
- 版本tag (v*): 永久保留
- 固定tag (latest, main, develop, master): 永久保留
- 缓存镜像 (*-cache): 永久保留
- 受保护tag (--protected-tags): 永久保留（如当前运行中镜像）
- PR预览tag (pr-*):
  - --pr-sha模式：删除指定PR commit的镜像（PR关闭时触发）
  - cron模式：通过Gitea API检查PR状态，已关闭/合并的删除
- 普通commit hash tag: 保留最近 N 个（默认20），老的删除

使用方式：
  # 预览（不实际删除）
  python3 acr_cleanup.py --dry-run

  # 实际执行（cron模式）
  python3 acr_cleanup.py --execute

  # 保留最近30个commit镜像
  python3 acr_cleanup.py --keep 30 --execute

  # PR关闭时清理指定commit的PR镜像
  python3 acr_cleanup.py --pr-sha abc123def --execute

  # 传入受保护tag列表（运行中镜像白名单）
  python3 acr_cleanup.py --protected-tags "sha1,sha2" --execute
"""

import argparse
import base64
import json
import os
import sys
import urllib.error
import urllib.request
from datetime import datetime, timedelta, timezone

# ========== 配置 ==========

REGISTRY = os.environ.get("ACR_REGISTRY", "xiaoxia-registry.cn-hangzhou.cr.aliyuncs.com")
AUTH_URL = "https://dockerauth.cn-hangzhou.aliyuncs.com/auth"
SERVICE = os.environ.get("ACR_SERVICE", "registry.aliyuncs.com:cn-hangzhou:china:cri-fvec8o9q4mmxrkaa")
NAMESPACE = os.environ.get("ACR_NAMESPACE", "xiaoxiakeji")
USERNAME = os.environ.get("ACR_USERNAME", "")
PASSWORD = os.environ.get("ACR_PASSWORD", "")

# Gitea配置（用于PR状态检查）
GITEA_URL = os.environ.get("GITEA_URL", "https://git.xiaoxiajianji.com")
GITEA_TOKEN = os.environ.get("GITEA_TOKEN", "")
GITEA_REPO = os.environ.get("GITEA_REPO", "xiaoxia/xiaoxia-saas")

REPOS = [
    "xiaoxia-saas-api",
    "xiaoxia-saas-worker",
    "xiaoxia-saas-web",
    "api-cache",
    "worker-cache",
    "web-cache",
]

# 缓存镜像仓库（所有tag永久保留）
CACHE_REPOS = {"api-cache", "worker-cache", "web-cache"}

# OCI / Docker manifest types
ACCEPT_INDEX = "application/vnd.oci.image.index.v1+json"
ACCEPT_MANIFEST_OCI = "application/vnd.oci.image.manifest.v1+json"
ACCEPT_MANIFEST_V2 = "application/vnd.docker.distribution.manifest.v2+json"


# ========== Registry API ==========


def get_token(repo, action="pull"):
    """获取仓库访问token"""
    scope = "repository:" + NAMESPACE + "/" + repo + ":" + action
    token_url = AUTH_URL + "?service=" + SERVICE + "&scope=" + scope
    req = urllib.request.Request(token_url)
    req.add_header("Authorization", "Basic " + base64.b64encode((USERNAME + ":" + PASSWORD).encode()).decode())
    with urllib.request.urlopen(req) as resp:
        data = json.loads(resp.read())
        return data.get("token", "")


def get_tags(repo, token):
    """获取仓库所有tag"""
    url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/tags/list?n=1000"
    req = urllib.request.Request(url)
    req.add_header("Authorization", "Bearer " + token)
    with urllib.request.urlopen(req) as resp:
        data = json.loads(resp.read())
        return data.get("tags", []) or []


def http_get_json(url, token, accept_header):
    """带Authorization的GET请求，返回(json_data, headers)"""
    req = urllib.request.Request(url)
    req.add_header("Authorization", "Bearer " + token)
    req.add_header("Accept", accept_header)
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read()), resp.headers


def get_manifest_info(repo, tag, token):
    """
    获取tag的manifest信息。
    返回: {digest, created, media_type, error}
    """
    url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/manifests/" + tag
    result = {"digest": "", "created": "", "media_type": "", "error": ""}

    # 先尝试 OCI index 格式
    try:
        data, headers = http_get_json(url, token, ACCEPT_INDEX)
        top_digest = headers.get("Docker-Content-Digest", "")
        result["digest"] = top_digest
        result["media_type"] = data.get("mediaType", ACCEPT_INDEX)

        manifests = data.get("manifests", [])
        amd64_manifest = None
        for m in manifests:
            arch = m.get("platform", {}).get("architecture", "")
            if arch == "amd64":
                amd64_manifest = m
                break
        if not amd64_manifest and manifests:
            amd64_manifest = manifests[0]

        if amd64_manifest:
            inner_digest = amd64_manifest["digest"]
            inner_url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/manifests/" + inner_digest
            try:
                inner_data, _ = http_get_json(inner_url, token, ACCEPT_MANIFEST_OCI)
            except Exception:
                inner_data, _ = http_get_json(inner_url, token, ACCEPT_MANIFEST_V2)

            config_digest = inner_data.get("config", {}).get("digest", "")
            if config_digest:
                blob_url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/blobs/" + config_digest
                try:
                    blob_data, _ = http_get_json(blob_url, token, "application/json")
                    result["created"] = blob_data.get("created", "")
                except Exception:
                    pass
        return result
    except urllib.error.HTTPError:
        pass

    # 再尝试普通 OCI manifest 格式
    try:
        data, headers = http_get_json(url, token, ACCEPT_MANIFEST_OCI)
        result["digest"] = headers.get("Docker-Content-Digest", "")
        result["media_type"] = data.get("mediaType", ACCEPT_MANIFEST_OCI)
        config_digest = data.get("config", {}).get("digest", "")
        if config_digest:
            blob_url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/blobs/" + config_digest
            try:
                blob_data, _ = http_get_json(blob_url, token, "application/json")
                result["created"] = blob_data.get("created", "")
            except Exception:
                pass
        return result
    except urllib.error.HTTPError:
        pass

    # 最后试 Docker v2 格式
    try:
        data, headers = http_get_json(url, token, ACCEPT_MANIFEST_V2)
        result["digest"] = headers.get("Docker-Content-Digest", "")
        result["media_type"] = data.get("mediaType", ACCEPT_MANIFEST_V2)
        config_digest = data.get("config", {}).get("digest", "")
        if config_digest:
            blob_url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/blobs/" + config_digest
            try:
                blob_data, _ = http_get_json(blob_url, token, "application/json")
                result["created"] = blob_data.get("created", "")
            except Exception:
                pass
        return result
    except urllib.error.HTTPError as e:
        result["error"] = "HTTP " + str(e.code) + " " + e.read().decode()[:200]

    return result


def delete_manifest(repo, digest, token):
    """按digest删除manifest（会级联删除所有指向它的tag）"""
    url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/manifests/" + digest
    req = urllib.request.Request(url, method="DELETE")
    req.add_header("Authorization", "Bearer " + token)
    req.add_header("Accept", ACCEPT_INDEX)
    req.add_header("Accept", ACCEPT_MANIFEST_OCI)
    req.add_header("Accept", ACCEPT_MANIFEST_V2)
    try:
        with urllib.request.urlopen(req) as resp:
            return True, resp.status
    except urllib.error.HTTPError as e:
        return False, str(e.code) + " " + e.read().decode()[:200]


# ========== Gitea API ==========


def gitea_get_open_prs():
    """获取所有打开的PR编号列表"""
    if not GITEA_TOKEN:
        print("  警告: 无GITEA_TOKEN，跳过PR状态检查")
        return None

    open_prs = set()
    page = 1
    while True:
        url = GITEA_URL + "/api/v1/repos/" + GITEA_REPO + "/pulls?state=open&page=" + str(page) + "&limit=50"
        req = urllib.request.Request(url)
        req.add_header("Authorization", "token " + GITEA_TOKEN)
        try:
            with urllib.request.urlopen(req) as resp:
                data = json.loads(resp.read())
                if not data:
                    break
                for pr in data:
                    open_prs.add(pr.get("number", 0))
                if len(data) < 50:
                    break
                page += 1
        except Exception as e:
            print(f"  警告: 获取Gitea PR列表失败: {e}")
            return None

    return open_prs


def gitea_get_pr_commits(pr_number):
    """获取指定PR的所有commit sha"""
    if not GITEA_TOKEN:
        return []

    url = GITEA_URL + "/api/v1/repos/" + GITEA_REPO + "/pulls/" + str(pr_number) + "/commits?limit=100"
    req = urllib.request.Request(url)
    req.add_header("Authorization", "token " + GITEA_TOKEN)
    try:
        with urllib.request.urlopen(req) as resp:
            data = json.loads(resp.read())
            return [c.get("sha", "") for c in data]
    except Exception as e:
        print(f"  警告: 获取PR #{pr_number} commits失败: {e}")
        return []


# ========== 工具函数 ==========


def parse_time(created_str):
    """解析ISO时间字符串"""
    if not created_str:
        return datetime.min.replace(tzinfo=timezone.utc)
    try:
        if created_str.endswith("Z"):
            created_str = created_str[:-1] + "+00:00"
        return datetime.fromisoformat(created_str)
    except Exception:
        return datetime.min.replace(tzinfo=timezone.utc)


def is_version_tag(tag):
    """判断是否是版本tag (v1.2.3, v0.1.0-alpha等)"""
    return tag.startswith("v") and len(tag) > 1 and tag[1].isdigit()


def is_fixed_tag(tag):
    """判断是否是固定tag"""
    return tag in ("latest", "main", "develop", "master", "dev", "stable")


def is_pr_tag(tag):
    """判断是否是PR预览tag (pr-<sha>)"""
    return tag.startswith("pr-")


def extract_sha_from_pr_tag(tag):
    """从pr-<sha> tag中提取sha"""
    if tag.startswith("pr-"):
        return tag[3:]
    return tag


def is_in_protected_list(tag, protected_set):
    """检查tag是否在受保护列表中"""
    if not protected_set:
        return False
    # 精确匹配
    if tag in protected_set:
        return True
    # 前缀匹配（commit hash可能是完整或短的）
    for p in protected_set:
        if tag.startswith(p) or p.startswith(tag):
            return True
    return False


# ========== 核心清理逻辑 ==========


def cleanup_repo(repo, keep_count, dry_run, protected_tags, pr_sha=None, pr_open_set=None):
    """
    清理单个仓库

    Args:
        repo: 仓库名
        keep_count: 保留最近N个commit tag
        dry_run: 是否预览模式
        protected_tags: 受保护tag集合（白名单）
        pr_sha: 指定PR commit sha（PR关闭模式），None表示cron模式
        pr_open_set: 打开的PR编号集合（cron模式用）

    Returns:
        (总tag数, 删除数)
    """
    print("=" * 60)
    print("仓库:", repo)
    print("=" * 60)

    # 缓存仓库不清理
    if repo in CACHE_REPOS:
        token_pull = get_token(repo, "pull")
        tags = get_tags(repo, token_pull)
        print("  缓存仓库，跳过清理 (共", len(tags), "个tag)")
        return len(tags), 0

    token_pull = get_token(repo, "pull")
    tags = get_tags(repo, token_pull)
    print("  总tag数:", len(tags))

    if not tags:
        print("  无tag，跳过")
        return 0, 0

    # ========== PR-SHA模式：只删除指定commit的PR镜像 ==========
    if pr_sha:
        pr_tags_to_del = [
            t
            for t in tags
            if t.startswith("pr-" + pr_sha) or t == "pr-" + pr_sha or pr_sha.startswith(extract_sha_from_pr_tag(t))
        ]
        if not pr_tags_to_del:
            print(f"  未找到PR镜像: pr-{pr_sha[:12]}")
            return len(tags), 0

        print(f"  找到 {len(pr_tags_to_del)} 个PR镜像待删除:")
        for t in pr_tags_to_del:
            print(f"    - {t}")

        to_delete = []
        for tag in pr_tags_to_del:
            info = get_manifest_info(repo, tag, token_pull)
            if info["digest"]:
                to_delete.append({"tag": tag, "digest": info["digest"], "created": info["created"]})
            else:
                print(f"    警告: {tag} 无法获取digest，跳过")

        return _execute_delete(repo, to_delete, dry_run, len(tags))

    # ========== Cron模式：全量清理 ==========

    # 分类
    version_tags = []
    fixed_tags = []
    pr_tags_list = []
    commit_tags = []

    for tag in tags:
        if is_version_tag(tag):
            version_tags.append(tag)
        elif is_fixed_tag(tag):
            fixed_tags.append(tag)
        elif is_pr_tag(tag):
            pr_tags_list.append(tag)
        else:
            commit_tags.append(tag)

    print("  版本tag (v*):", len(version_tags), "-> 永久保留")
    print("  固定tag:", len(fixed_tags), "-> 永久保留")
    print("  PR预览tag (pr-*):", len(pr_tags_list), "-> 已关闭PR的删除")
    print("  Commit hash tag:", len(commit_tags), "-> 保留最近", keep_count, "个")
    print("  白名单tag:", len(protected_tags), "个")

    # --- PR tag清理：检查PR状态 ---
    pr_to_delete = []
    if pr_tags_list:
        print()
        print("  检查PR镜像状态...")

        # 策略：有Gitea token则检查PR状态，否则按时间保留7天
        if pr_open_set is not None:
            # 通过Gitea API检查每个PR镜像对应的PR是否还开着
            # 注意：pr tag是pr-<sha>，sha可能属于某个PR
            # 简化策略：收集所有打开PR的commit sha，在白名单里的保留
            print("    模式: Gitea PR状态检查")
            open_pr_shas = set()
            # 这里做了简化：因为每个PR都查commits太慢，我们用另一种方式
            # 对于PR tag，先尝试匹配PR编号（如果tag名里有编号），否则按时间
            # 实际pr-<sha>没法直接知道PR编号，所以降级为按时间+打开PR的head sha白名单
            open_head_shas = set()
            page = 1
            while True:
                url = GITEA_URL + "/api/v1/repos/" + GITEA_REPO + "/pulls?state=open&page=" + str(page) + "&limit=50"
                req = urllib.request.Request(url)
                req.add_header("Authorization", "token " + GITEA_TOKEN)
                try:
                    with urllib.request.urlopen(req) as resp:
                        data = json.loads(resp.read())
                        if not data:
                            break
                        for pr in data:
                            head_sha = pr.get("head", {}).get("sha", "")
                            if head_sha:
                                open_head_shas.add(head_sha)
                                open_head_shas.add(head_sha[:7])
                                open_head_shas.add(head_sha[:12])
                        if len(data) < 50:
                            break
                        page += 1
                except Exception:
                    break

            deleted_count = 0
            for tag in pr_tags_list:
                sha = extract_sha_from_pr_tag(tag)
                # 检查是否是打开PR的head sha
                is_open_pr = False
                for ohs in open_head_shas:
                    if sha.startswith(ohs) or ohs.startswith(sha):
                        is_open_pr = True
                        break
                if not is_open_pr:
                    info = get_manifest_info(repo, tag, token_pull)
                    if info["digest"]:
                        pr_to_delete.append({"tag": tag, "digest": info["digest"], "created": info["created"]})
                        deleted_count += 1
            print(f"    打开PR数: {len(open_head_shas)}个head sha")
            print(f"    将删除PR镜像: {deleted_count}个")
        else:
            # 无Gitea token，降级为按7天保留
            print("    模式: 按时间保留7天（无Gitea token降级）")
            cutoff = datetime.now(timezone.utc) - timedelta(days=7)
            for tag in pr_tags_list:
                info = get_manifest_info(repo, tag, token_pull)
                created = parse_time(info["created"])
                if created < cutoff and info["digest"]:
                    pr_to_delete.append({"tag": tag, "digest": info["digest"], "created": info["created"]})
            print(f"    将删除PR镜像: {len(pr_to_delete)}个")

    # --- Commit tag清理：保留最近N个 ---
    print()
    print("  获取commit tag创建时间...")
    commit_tag_infos = []
    errors = 0
    for i, tag in enumerate(commit_tags):
        info = get_manifest_info(repo, tag, token_pull)
        if info["error"] or not info["digest"]:
            errors += 1
        commit_tag_infos.append({"tag": tag, "digest": info["digest"], "created": info["created"]})
        if (i + 1) % 20 == 0:
            print("    已获取", i + 1, "/", len(commit_tags), "...")

    if errors:
        print("    注意:", errors, "个tag获取manifest失败")

    # 按时间倒序排序
    commit_tag_infos.sort(key=lambda x: parse_time(x["created"]), reverse=True)

    # 确定要删除的commit tag
    commit_to_delete = []
    if len(commit_tag_infos) > keep_count:
        commit_to_delete = commit_tag_infos[keep_count:]
        print(f"  保留前{keep_count}个commit tag，删除{len(commit_to_delete)}个")

        # 白名单过滤：受保护的tag不删除
        if protected_tags:
            before = len(commit_to_delete)
            commit_to_delete = [t for t in commit_to_delete if not is_in_protected_list(t["tag"], protected_tags)]
            removed = before - len(commit_to_delete)
            if removed > 0:
                print(f"  白名单保护: 跳过{removed}个运行中镜像")

        # 过滤无digest的
        commit_to_delete = [t for t in commit_to_delete if t["digest"]]
        print(f"  可删除(有digest): {len(commit_to_delete)}个")
    else:
        print(f"  commit tag数量不足{keep_count}个，无需清理")

    # --- 合并所有待删除项 ---
    all_to_delete = commit_to_delete + pr_to_delete

    # 再次过滤白名单（PR镜像也受白名单保护）
    if protected_tags:
        before = len(all_to_delete)
        all_to_delete = [t for t in all_to_delete if not is_in_protected_list(t["tag"], protected_tags)]
        removed = before - len(all_to_delete)
        if removed > 0:
            print(f"  白名单保护(PR镜像): 跳过{removed}个")

    return _execute_delete(repo, all_to_delete, dry_run, len(tags))


def _execute_delete(repo, to_delete, dry_run, total_tags):
    """执行删除操作"""
    if not to_delete:
        print()
        print("  无需删除任何tag")
        return total_tags, 0

    # 按digest去重
    seen_digests = set()
    unique_delete = []
    for item in to_delete:
        if item["digest"] and item["digest"] not in seen_digests:
            seen_digests.add(item["digest"])
            unique_delete.append(item)

    print()
    if dry_run:
        print(f"  [DRY RUN] 将删除{len(unique_delete)}个manifest（预览模式）")
        for item in unique_delete[:5]:
            created_str = item.get("created", "")[:10] or "未知"
            print(f"    - {item['tag'][:30]} ({created_str})")
        if len(unique_delete) > 5:
            print(f"    ... 还有{len(unique_delete) - 5}个")
        return total_tags, len(unique_delete)

    token_delete = get_token(repo, "delete")
    deleted = 0
    failed = 0

    print(f"  开始删除{len(unique_delete)}个唯一manifest...")
    for item in unique_delete:
        success, result = delete_manifest(repo, item["digest"], token_delete)
        if success:
            deleted += 1
            print(f"    已删除: {item['tag'][:30]}")
        else:
            failed += 1
            print(f"    删除失败: {item['tag'][:30]} - {result}")

    print()
    print(f"  删除完成: 成功{deleted}个，失败{failed}个")
    return total_tags, deleted


# ========== 主函数 ==========


def main():
    parser = argparse.ArgumentParser(description="ACR镜像清理工具（增强版）")
    parser.add_argument("--keep", type=int, default=20, help="保留最近N个commit hash tag（默认20）")
    parser.add_argument("--dry-run", action="store_true", help="预览模式，不实际删除")
    parser.add_argument("--execute", action="store_true", help="实际执行删除")
    parser.add_argument("--repo", type=str, default="", help="只清理指定仓库")
    parser.add_argument("--pr-sha", type=str, default="", help="PR关闭模式：删除指定commit sha的PR镜像")
    parser.add_argument("--protected-tags", type=str, default="", help="受保护tag列表，逗号分隔（运行中镜像白名单）")
    parser.add_argument("--skip-pr-check", action="store_true", help="跳过Gitea PR状态检查（纯按时间清理PR镜像）")
    args = parser.parse_args()

    # 必须指定 --dry-run 或 --execute
    if not args.dry_run and not args.execute:
        print("请指定 --dry-run（预览）或 --execute（执行）")
        print()
        print("示例:")
        print("  python3 acr_cleanup.py --dry-run           # 预览清理效果")
        print("  python3 acr_cleanup.py --execute           # 实际执行清理")
        print("  python3 acr_cleanup.py --pr-sha abc123 --execute  # PR关闭时清理")
        sys.exit(1)

    # 凭证检查
    global USERNAME, PASSWORD
    if not USERNAME or not PASSWORD:
        try:
            docker_config_path = os.path.expanduser("~/.docker/config.json")
            with open(docker_config_path) as f:
                config = json.load(f)
            auth = config.get("auths", {}).get(REGISTRY, {}).get("auth", "")
            if auth:
                creds = base64.b64decode(auth).decode().strip()
                USERNAME, PASSWORD = creds.split(":", 1)
        except Exception:
            pass

    if not USERNAME or not PASSWORD:
        print("错误: 缺少ACR凭证，请设置 ACR_USERNAME 和 ACR_PASSWORD 环境变量")
        print("或确保已执行 docker login", REGISTRY)
        sys.exit(1)

    # 解析受保护tag
    protected_tags = set()
    if args.protected_tags:
        protected_tags = set(t.strip() for t in args.protected_tags.split(",") if t.strip())

    dry_run = args.dry_run or not args.execute
    mode = "预览模式" if dry_run else "执行模式"

    print("=" * 60)
    print("ACR 镜像清理工具（增强版）-", mode)
    print("=" * 60)
    print("Registry:", REGISTRY)
    print("Namespace:", NAMESPACE)
    if args.pr_sha:
        print("模式: PR关闭清理")
        print("PR commit SHA:", args.pr_sha[:12])
    else:
        print("模式: Cron全量清理")
        print("保留commit tag数:", args.keep)
        print("PR状态检查:", "关闭" if args.skip_pr_check else "开启")
    if protected_tags:
        print("白名单tag数:", len(protected_tags))
    print()

    # PR模式不需要查Gitea
    pr_open_set = None
    if not args.pr_sha and not args.skip_pr_check and GITEA_TOKEN:
        print("获取打开的PR列表...")
        pr_open_set = gitea_get_open_prs()
        if pr_open_set is not None:
            print(f"  打开的PR: {len(pr_open_set)}个")
        print()

    repos_to_clean = REPOS
    if args.repo:
        repos_to_clean = [args.repo]

    total_deleted = 0
    total_tags = 0
    for repo in repos_to_clean:
        count, deleted = cleanup_repo(
            repo, args.keep, dry_run, protected_tags, pr_sha=args.pr_sha, pr_open_set=pr_open_set
        )
        total_tags += count
        total_deleted += deleted
        print()

    print("=" * 60)
    print("清理完成")
    print("  总tag数:", total_tags)
    if dry_run:
        print("  预览将删除(去重后):", total_deleted, "个manifest")
    else:
        print("  已删除:", total_deleted, "个manifest")
    print("=" * 60)


if __name__ == "__main__":
    main()