vendor_magisk/magisk.py

#!/usr/bin/env python3
import argparse
import hashlib
import os
import re
import shutil
import subprocess
import urllib.request
import zipfile


def calculate_md5(file_path):
    """计算文件的 MD5 值"""
    hash_md5 = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def download_file(url, file_path):
    """下载文件"""
    print(f"==> Downloading from {url}")
    urllib.request.urlretrieve(url, file_path)


def run_command(cmd):
    """运行命令"""
    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Command failed: {' '.join(cmd)}")
        raise e


class Magisk:
    def __init__(self, target_arch="arm64"):
        self.download_loc = os.path.join(os.path.dirname(os.path.abspath(__file__)), "downloads")
        self.dl_link = "https://github.com/topjohnwu/Magisk/releases/download/v30.2/Magisk-v30.2.apk"
        self.dl_file_name = os.path.join(self.download_loc, "magisk.apk")
        self.act_md5 = "2691c30ccf059af2536cb0af803c787c"
        self.extract_to = os.path.join(os.path.dirname(os.path.abspath(__file__)), "temp")
        self.copy_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "magisk")
        self.target_arch = target_arch
        
        # 架构映射
        self.arch_map = {
            "x86": "x86",
            "x86_64": "x86_64",
            "arm": "armeabi-v7a",
            "arm64": "arm64-v8a"
        }

    def download(self):
        """下载 Magisk APK"""
        print("==> Downloading latest Magisk now .....")
        
        # 确保下载目录存在
        os.makedirs(self.download_loc, exist_ok=True)
        
        # 检查并下载 Magisk APK
        need_download = True
        if os.path.exists(self.dl_file_name):
            print("==> Checking existing Magisk APK...")
            actual_md5 = calculate_md5(self.dl_file_name)
            if actual_md5 == self.act_md5:
                print("==> Magisk APK already exists and verified!")
                need_download = False
            else:
                print(f"==> MD5 mismatch. Expected: {self.act_md5}, Got: {actual_md5}")
        
        if need_download:
            download_file(self.dl_link, self.dl_file_name)
            actual_md5 = calculate_md5(self.dl_file_name)
            if actual_md5 != self.act_md5:
                raise Exception(f"Downloaded file MD5 verification failed! Expected: {self.act_md5}, Got: {actual_md5}")
            print("==> Download completed and verified!")

    def extract(self):
        """解压 APK 文件"""
        print("==> Extracting Magisk APK...")
        
        # 清理并创建解压目录
        shutil.rmtree(self.extract_to, ignore_errors=True)
        os.makedirs(self.extract_to, exist_ok=True)
        
        # 解压 APK
        with zipfile.ZipFile(self.dl_file_name) as z:
            z.extractall(self.extract_to)

    def copy(self):
        """复制文件到目标目录"""
        print("==> Copying magisk files now ...")
        
        # 清理并创建目标目录
        if os.path.exists(self.copy_dir):
            shutil.rmtree(self.copy_dir)
        os.makedirs(self.copy_dir, exist_ok=True)
        
        # 复制主要架构的库文件
        lib_dir = os.path.join(self.extract_to, "lib", self.arch_map[self.target_arch])
        if os.path.exists(lib_dir):
            for parent, dirnames, filenames in os.walk(lib_dir):
                for filename in filenames:
                    o_path = os.path.join(lib_dir, filename)
                    so_name = re.search(r'lib(.*)\.so', filename)
                    if so_name:
                        n_path = os.path.join(self.copy_dir, so_name.group(1))
                        shutil.copyfile(o_path, n_path)
                        run_command(["chmod", "+x", n_path])
        
        # 复制 arm32 的 magisk 二进制文件（如果存在）
        lib32_path = os.path.join(self.extract_to, "lib", "armeabi-v7a")
        magisk32_src = os.path.join(lib32_path, "libmagisk32.so")
        magisk32_dst = os.path.join(self.copy_dir, "magisk32")
        if os.path.exists(magisk32_src):
            shutil.copyfile(magisk32_src, magisk32_dst)
            run_command(["chmod", "+x", magisk32_dst])
        
        # 复制 magisk.apk 到目标目录
        apk_dst = os.path.join(os.path.dirname(os.path.abspath(__file__)), "magisk.apk")
        if os.path.exists(self.dl_file_name):
            shutil.copyfile(self.dl_file_name, apk_dst)
        
        print("==> Magisk installation completed!")

    def cleanup(self):
        """清理临时文件和目录"""
        print("==> Cleaning up temporary files...")
        
        # 清理解压临时目录
        if os.path.exists(self.extract_to):
            shutil.rmtree(self.extract_to, ignore_errors=True)
            print(f"==> Removed temporary extraction directory: {self.extract_to}")
        
        # 清理下载目录
        if os.path.exists(self.download_loc):
            shutil.rmtree(self.download_loc, ignore_errors=True)
            print(f"==> Removed download directory: {self.download_loc}")

    def run(self):
        """执行完整的构建流程"""
        try:
            self.download()
            self.extract()
            self.copy()
            success = True
        except Exception as e:
            print(f"Error: {e}")
            success = False
        finally:
            self.cleanup()
        
        return success


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='Magisk vendor package builder')
    
    parser.add_argument('--arch', 
                        choices=['x86', 'x86_64', 'arm', 'arm64'],
                        default='arm64',
                        help='指定目标架构 (默认: arm64)')
    
    args = parser.parse_args()
    
    print(f"==> Using target architecture: {args.arch}")
    
    # 创建 Magisk 实例
    magisk = Magisk(target_arch=args.arch)
    
    # 执行构建
    success = magisk.run()
    
    if success:
        print("==> Magisk vendor package created successfully!")
        print("==> Temporary files cleaned up")
    else:
        print("==> Failed to create Magisk vendor package!")
        exit(1)


if __name__ == '__main__':
    main()