xiaoxia-saas/scripts/ci/acr_cleanup.py

#!/usr/bin/env python3
"""
ACR 镜像清理脚本
策略：
- 版本tag (v*): 永久保留
- 固定tag (latest, main, develop, master): 永久保留
- 缓存镜像 (*-cache): 永久保留
- PR预览tag (pr-*): 保留 N 天（默认7天）
- 普通commit hash tag: 保留最近 N 个（默认20），老的删除

使用方式：
  python3 acr_cleanup.py --dry-run          # 预览，不实际删除
  python3 acr_cleanup.py --execute          # 实际执行删除
  python3 acr_cleanup.py --keep 20 --execute # 保留最近20个
"""

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", "")

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"


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信息，支持OCI index和普通manifest两种格式。
    返回: {digest, created, media_type}
    - digest: 顶层manifest的digest（用于删除）
    - created: 镜像创建时间
    """
    url = "https://" + REGISTRY + "/v2/" + NAMESPACE + "/" + repo + "/manifests/" + tag
    result = {"digest": "", "created": "", "media_type": "", "error": ""}

    # 先尝试 OCI index 格式（ACR多用这种）
    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)

        # OCI index：找amd64的manifest，再取config blob
        manifests = data.get("manifests", [])
        amd64_manifest = None
        for m in manifests:
            arch = m.get("platform", {}).get("architecture", "")
            if arch == "amd64":
                amd64_manifest = m
                break
        # 没有amd64就用第一个
        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:
                # 退而求其次用v2格式
                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]


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"""
    return tag.startswith("pr-")


def cleanup_repo(repo, keep_count, pr_days, dry_run):
    """清理单个仓库"""
    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))

    # 分类
    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_days, "天")
    print("  Commit hash tag:", len(commit_tags), "-> 保留最近", keep_count, "个")

    # 获取所有commit tag的创建时间
    print()
    print("  获取commit tag创建时间...")
    tag_info_list = []
    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
            # 取不到信息的tag，放到最后（最旧处理），但标记一下
            tag_info_list.append({"tag": tag, "digest": info["digest"], "created": "", "error": info.get("error", "")})
        else:
            tag_info_list.append({"tag": tag, "digest": info["digest"], "created": info["created"], "error": ""})
        if (i + 1) % 20 == 0:
            print("    已获取", i + 1, "/", len(commit_tags), "...")

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

    # 按时间倒序排序（空时间放最后）
    tag_info_list.sort(key=lambda x: parse_time(x["created"]), reverse=True)

    # 确定要删除的commit tag
    to_delete = []
    if len(tag_info_list) > keep_count:
        to_delete = tag_info_list[keep_count:]
        print("  保留前", keep_count, "个commit tag，删除", len(to_delete), "个")
        # 打印保留范围
        kept = tag_info_list[:keep_count]
        valid_kept = [t for t in kept if t["created"]]
        if valid_kept:
            print("  最早保留:", valid_kept[-1]["tag"][:12], "(" + valid_kept[-1]["created"][:10] + ")")
        # 保护当前构建的tag（通过PROTECTED_TAG环境变量传入，如GITHUB_SHA）
        protected_tag = os.environ.get("PROTECTED_TAG", "").strip()
        if protected_tag:
            before = len(to_delete)
            to_delete = [t for t in to_delete if not t["tag"].startswith(protected_tag)]
            removed = before - len(to_delete)
            if removed > 0:
                print(f"  保护当前构建tag: {protected_tag[:12]} (跳过{removed}个)")

        to_del_valid = [t for t in to_delete if t["digest"]]
        print("  可删除(有digest):", len(to_del_valid), "个")
    else:
        print("  commit tag数量不足", keep_count, "，无需清理")

    # PR tag按时间清理
    pr_to_delete = []
    if pr_tags_list:
        cutoff = datetime.now(timezone.utc) - timedelta(days=pr_days)
        print()
        print("  检查PR预览tag（超过", pr_days, "天删除）...")
        for tag in pr_tags_list:
            info = get_manifest_info(repo, tag, token_pull)
            created = parse_time(info["created"])
            if created < cutoff:
                pr_to_delete.append({"tag": tag, "digest": info["digest"], "created": info["created"]})
        print("  PR tag将删除:", len(pr_to_delete), "个")

    all_to_delete = [t for t in to_delete if t["digest"]] + [t for t in pr_to_delete if t["digest"]]

    if not all_to_delete:
        print()
        print("  无需删除任何tag")
        return len(tags), 0

    # 执行删除
    print()
    if dry_run:
        print("  [DRY RUN] 将删除", len(all_to_delete), "个tag（预览模式，不实际删除）")
        # 去重digest
        unique_digests = set(t["digest"] for t in all_to_delete if t["digest"])
        print("  去重后唯一digest数:", len(unique_digests))
        for item in all_to_delete[:5]:
            created_str = item.get("created", "")[:10] or "未知"
            print("    -", item["tag"][:20], "(" + created_str + ")")
        if len(all_to_delete) > 5:
            print("    ... 还有", len(all_to_delete) - 5, "个")
        return len(tags), len(unique_digests)

    token_delete = get_token(repo, "delete")
    deleted = 0
    failed = 0
    # 按digest去重，避免重复删除同一镜像
    seen_digests = set()
    unique_delete = []
    for item in all_to_delete:
        if item["digest"] and item["digest"] not in seen_digests:
            seen_digests.add(item["digest"])
            unique_delete.append(item)

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

    print()
    print("  删除完成: 成功", deleted, "个，失败", failed, "个")
    return len(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("--pr-days", type=int, default=7, help="PR预览tag保留天数（默认7天）")
    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="只清理指定仓库")
    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 --keep 20 --execute # 保留最近20个")
        sys.exit(1)

    # 凭证检查
    global USERNAME, PASSWORD
    if not USERNAME or not PASSWORD:
        # 尝试从docker config读取
        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)

    dry_run = args.dry_run or not args.execute
    mode = "预览模式" if dry_run else "执行模式"
    print("ACR 镜像清理工具 -", mode)
    print("Registry:", REGISTRY)
    print("Namespace:", NAMESPACE)
    print("保留commit tag数:", args.keep)
    print("PR预览保留天数:", args.pr_days)
    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, args.pr_days, dry_run)
        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()