omegafox/scripts/package.py

#!/usr/bin/env python3

import argparse
import glob
import os
import shutil
import sys
import tempfile
from shlex import join

from _mixin import find_src_dir, get_moz_target, list_files, run, temp_cd

UNNEEDED_PATHS = {'uninstall', 'pingsender.exe', 'pingsender', 'vaapitest', 'glxtest'}


def add_includes_to_package(package_file, includes, fonts, new_file, target):
    with tempfile.TemporaryDirectory() as temp_dir:
        # Extract package
        run(join(['7z', 'x', package_file, f'-o{temp_dir}']), exit_on_fail=False)
        # Delete package_file
        os.remove(package_file)
        if package_file.endswith('.tar.xz'):
            # Rerun on the tar file
            package_tar = package_file[:-3]  # Remove ".xz"
            return add_includes_to_package(
                package_file=os.path.join(temp_dir, package_tar),
                includes=includes,
                fonts=fonts,
                new_file=new_file,
                target=target,
            )

        # Move contents out of testfox folder if it exists
        old_testfox_dir = os.path.join(temp_dir, 'testfox')
        testfox_dir = os.path.join(temp_dir, 'testfox-folder')
        if os.path.exists(old_testfox_dir):
            # Rename testfox_dir
            os.rename(old_testfox_dir, testfox_dir)
            for item in os.listdir(testfox_dir):
                shutil.move(os.path.join(testfox_dir, item), temp_dir)
            os.rmdir(testfox_dir)

        # Create target_dir
        target_dir = temp_dir

        # Add includes
        for include in includes or []:
            if os.path.exists(include):
                if os.path.isdir(include):
                    shutil.copytree(
                        include,
                        os.path.join(target_dir, os.path.basename(include)),
                        dirs_exist_ok=True,
                    )
                else:
                    shutil.copy2(include, target_dir)

        # Add the font folders under fonts/
        fonts_dir = os.path.join(target_dir, 'fonts')
        for font in fonts or []:
            shutil.copytree(
                os.path.join('bundle', 'fonts', font),
                os.path.join(fonts_dir, font),
                dirs_exist_ok=True,
            )

        # Add hardened-malloc
        shutil.copy2(
            os.path.join('libhadened_malloc.so'),
            os.path.join(target_dir, 'libhardened_malloc.so'),
            dirs_exist_ok=True,
        )

        # Remove unneeded paths
        for path in UNNEEDED_PATHS:
            if os.path.isdir(os.path.join(target_dir, path)):
                shutil.rmtree(os.path.join(target_dir, path), ignore_errors=True)
            elif os.path.exists(os.path.join(target_dir, path)):
                os.remove(os.path.join(target_dir, path))

        # Update package
        run(join(['7z', 'u', new_file, f'{temp_dir}/*', '-r', '-mx=9']))


def get_args():
    """Get CLI parameters"""
    parser = argparse.ArgumentParser(
        description='Package Omegafox for different operating systems.'
    )
    parser.add_argument('os', choices=['linux'], help='Target operating system')
    parser.add_argument(
        '--includes', nargs='+', help='List of files or directories to include in the package'
    )
    parser.add_argument('--version', required=True, help='Omegafox version')
    parser.add_argument('--release', required=True, help='Omegafox release number')
    parser.add_argument(
        '--arch', choices=['x86_64', 'i686', 'arm64'], help='Architecture for Windows build'
    )
    parser.add_argument('--fonts', nargs='+', help='Font directories to include under fonts/')
    return parser.parse_args()


def main():
    """The main packaging function"""
    args = get_args()

    # Determine file extension based on OS
    file_extensions = {'linux': 'tar.xz', 'macos': 'dmg', 'windows': 'zip'}
    file_ext = file_extensions[args.os]

    # Build the package
    src_dir = find_src_dir('.', args.version, args.release)
    moz_target = get_moz_target(target=args.os, arch=args.arch)
    with temp_cd(src_dir):
        # Create package files
        run('./mach package')
        # Find package files
        search_path = os.path.abspath(
            f'obj-{moz_target}/dist/omegafox-{args.version}-{args.release}.*.{file_ext}'
        )

    # Copy package files
    for file in glob.glob(search_path):
        if 'xpt_artifacts' in file:
            print(f'Skipping xpt artifacts: {file}')
            continue
        print(f'Found package: {file}')
        # Copy to root
        shutil.copy2(file, '.')
        break
    else:
        print(f"Error: No package file found matching pattern: {search_path}")
        sys.exit(1)

    # Find the package file
    package_pattern = f'omegafox-{args.version}-{args.release}.en-US.*.{file_ext}'
    package_files = glob.glob(package_pattern)
    if not package_files:
        print(f"Error: No package file found matching pattern: {package_pattern}")
        exit(1)
    package_file = package_files[0]

    # Add includes to the package
    new_name = f'omegafox-{args.version}-{args.release}-{args.os[:3]}.{args.arch}.zip'
    add_includes_to_package(
        package_file=package_file,
        includes=args.includes,
        fonts=args.fonts,
        new_file=new_name,
        target=args.os,
    )

    print(f"Packaging complete for {args.os}")


if __name__ == '__main__':
    main()