import UnityPy
from PIL import Image


def convert_ab_to_png(assetbundle_path, output_png_path):
    env = UnityPy.load(assetbundle_path)

    for obj in env.objects:
        if obj.type == 28:  # ClassID 28 = Texture2D
            texture2d = obj.read()
            image = texture2d.image
            image.save(output_png_path)
            print(f"已提取 Texture2D 并保存到 {output_png_path}")
            return

    raise Exception("AssetBundle 中没有找到 Texture2D。")


# 示例用法
if __name__ == "__main__":
    import sys

    if len(sys.argv) != 3:
        print("用法: python ab2png.py 输入.ab 输出.png")
    else:
        convert_ab_to_png(sys.argv[1], sys.argv[2])