Nugget/tweaks/posterboard_tweak.py

from .tweak_classes import Tweak
from Sparserestore.restore import FileToRestore
import os
import zipfile

class PosterboardTweak(Tweak):
    def __init__(self):
        super().__init__(key=None)
        self.zip_path = None
        self.bundle_id = "com.apple.PosterBoard"
        self.resetting = False

    def recursive_add(self, files_to_restore: list[FileToRestore], curr_path: str, restore_path: str = "", isAdding: bool = False):
        for folder in sorted(os.listdir(curr_path)):
            if folder.startswith('.') or folder == "__MACOSX":
                continue
            if isAdding:
                # if file then add it, otherwise recursively call again
                if os.path.isfile(os.path.join(curr_path, folder)):
                    try:
                        files_to_restore.append(FileToRestore(
                            contents=None,
                            contents_path=os.path.join(curr_path, folder),
                            restore_path=f"{restore_path}/{folder}",
                            domain=f"AppDomain-{self.bundle_id}"
                        ))
                    except IOError:
                        print(f"Failed to open file: {folder}") # TODO: Add QDebug equivalent
                else:
                    self.recursive_add(files_to_restore, os.path.join(curr_path, folder), f"{restore_path}/{folder}", isAdding)
            else:
                # look for contents folder
                if folder == "Container":
                    self.recursive_add(files_to_restore, os.path.join(curr_path, folder), restore_path="/", isAdding=True)
                    return
                else:
                    self.recursive_add(files_to_restore, os.path.join(curr_path, folder), isAdding=False)

    def apply_tweak(self, files_to_restore: list[FileToRestore], output_dir: str):
        # unzip the file
        if not self.enabled:
            return
        if self.resetting:
            # null out the prb folder
            files_to_restore.append(FileToRestore(
                contents=b"",
                restore_path="/Library/Application Support/PRBPosterExtensionDataStore",
                domain=f"AppDomain-{self.bundle_id}"
            ))
            return
        elif self.zip_path == None:
            return
        with zipfile.ZipFile(self.zip_path, 'r') as zip_ref:
            zip_ref.extractall(output_dir)
        # add the files
        self.recursive_add(files_to_restore, curr_path=output_dir)