environment.py

import io
import os
import ntpath
import re
from typing import List, Callable, Dict, Union
from zipfile import ZipFile

from fsspec import AbstractFileSystem
from fsspec.implementations.local import LocalFileSystem


from .files import File, ObjectReader, SerializedFile
from .enums import FileType
from .helpers import ImportHelper
from .streams import EndianBinaryReader

reSplit = re.compile(r"(.*?([^\/\\]+?))\.split\d+")


class Environment:
    files: dict
    cabs: dict
    path: str
    local_files: List[str]
    local_files_simple: List[str]

    def __init__(self, *args, fs: AbstractFileSystem = None):
        self.files = {}
        self.cabs = {}
        self.path = None
        self.fs = fs or LocalFileSystem()
        self.local_files = []
        self.local_files_simple = []

        if args:
            for arg in args:
                if isinstance(arg, str):
                    if self.fs.isfile(arg):
                        if ntpath.splitext(arg)[-1] in [".apk", ".zip"]:
                            self.load_zip_file(arg)
                        else:
                            self.path = ntpath.dirname(arg)
                            if reSplit.match(arg):
                                self.load_files([arg])
                            else:
                                self.load_file(arg)
                    elif self.fs.isdir(arg):
                        self.path = arg
                        self.load_folder(arg)
                else:
                    self.path = None
                    self.load_file(file=arg)

        if len(self.files) == 1:
            self.file = list(self.files.values())[0]

        if self.path == "":
            self.path = os.getcwd()

    def load_files(self, files: List[str]):
        """Loads all files (list) into the Environment and merges .split files for common usage."""
        self.load_assets(files, lambda x: open(x, "rb"))

    def load_folder(self, path: str):
        """Loads all files in the given path and its subdirs into the Environment."""
        self.load_files(
            [
                self.fs.sep.join([root, f])
                for root, dirs, files in self.fs.walk(path)
                for f in files
            ]
        )

    def load(self, files: list):
        """Loads all files into the Environment."""
        self.files.update(
            {
                ntpath.basename(f): self.load_file(self.fs.open(f, "rb"), self, f)
                for f in files
                if self.fs.exists(f)
            }
        )

    def load_file(
        self,
        file: Union[io.IOBase, str],
        parent: Union["Environment", File] = None,
        name: str = None,
        is_dependency: bool = False,
    ):
        if not parent:
            parent = self

        if isinstance(file, str):
            split_match = reSplit.match(file)
            if split_match:
                basepath, basename = split_match.groups()
                file = []
                for i in range(0, 999):
                    item = f"{basepath}.split{i}"
                    if self.fs.exists(item):
                        with self.fs.open(item, "rb") as f:
                            file.append(f.read())
                    elif i:
                        break
                name = basepath
                file = b"".join(file)
            else:
                name = file
                if not os.path.exists(file):
                    # relative paths are in the asset directory, not the cwd
                    if not os.path.isabs(file):
                        file = os.path.join(self.path, file)
                    # Unity paths are case insensitive, so we need to find "Resources/Foo.asset" when the record says "resources/foo.asset"
                    if not os.path.exists(file):
                        file = ImportHelper.find_sensitive_path(self.path, file)
                    # nonexistent files might be packaging errors or references to Unity's global Library/
                    if file is None:
                        return
                file = self.fs.open(file, "rb")

        typ, reader = ImportHelper.check_file_type(file)

        stream_name = (
            name
            if name
            else getattr(
                file,
                "name",
                str(file.__hash__()) if hasattr(file, "__hash__") else "",
            )
        )

        if typ == FileType.ZIP:
            f = self.load_zip_file(file)
        else:
            f = ImportHelper.parse_file(
                reader, self, name=stream_name, typ=typ, is_dependency=is_dependency
            )
        
        if isinstance(f, (SerializedFile, EndianBinaryReader)):
            self.register_cab(stream_name, f)

        self.files[stream_name] = f


    def load_zip_file(self, value):
        buffer = None
        if isinstance(value, str) and self.fs.exists(value):
            buffer = open(value, "rb")
        elif isinstance(value, (bytes, bytearray)):
            buffer = io.BytesIO(value)
        elif isinstance(value, (io.BufferedReader, io.BufferedIOBase)):
            buffer = value

        z = ZipFile(buffer)
        self.load_assets(z.namelist(), lambda x: z.open(x, "r"))
        z.close()

    def save(self, pack="none", out_path="output"):
        """Saves all changed assets.
        Mark assets as changed using `.mark_changed()`.
        pack = "none" (default) or "lz4"
        """
        for fname, fitem in self.files.items():
            if getattr(fitem, "is_changed", False):
                with open(
                    self.fs.sep.join([out_path, ntpath.basename(fname)]), "wb"
                ) as out:
                    out.write(fitem.save(packer=pack))

    @property
    def objects(self) -> List[ObjectReader]:
        """Returns a list of all objects in the Environment."""

        def search(item):
            ret = []
            if not isinstance(item, Environment) and getattr(item, "objects", None):
                # serialized file
                if getattr(item, "is_dependency", False):
                    return []
                return [val for val in item.objects.values()]

            elif getattr(item, "files", None):  # WebBundle and BundleFile
                # bundle
                for item in item.files.values():
                    ret.extend(search(item))
                return ret

            return ret

        return search(self)

    @property
    def container(self) -> Dict[str, ObjectReader]:
        """Returns a dictionary of all objects in the Environment."""
        return {
            path: obj
            for f in self.files.values()
            if isinstance(f, File) and not f.is_dependency
            for path, obj in f.container.items()
        }

    @property
    def assets(self) -> list:
        """
        Lists all assets / SerializedFiles within this environment.
        """

        def gen_all_asset_files(file, ret=[]):
            for f in getattr(file, "files", {}).values():
                if getattr(f, "is_dependency", False):
                    continue
                if isinstance(f, SerializedFile):
                    ret.append(f)
                else:
                    gen_all_asset_files(f, ret)
            return ret

        return gen_all_asset_files(self)

    def get(self, key: str, default=None):
        return getattr(self, key, default)

    def register_cab(self, name: str, item: File) -> None:
        """
        Registers a cab file.

        Parameters
        ----------
        name : str
            The name of the cab file.
        item : File
            The file to register.
        """
        self.cabs[simplify_name(name)] = item

    def get_cab(self, name: str) -> File:
        """
        Returns the cab file with the given name.

        Parameters
        ----------
        name : str
            The name of the cab file.

        Returns
        -------
        File
            The cab file.
        """
        return self.cabs.get(simplify_name(name), None)

    def load_assets(self, assets: List[str], open_f: Callable[[str], io.IOBase]):
        """
        Load all assets from a list of files via the given open_f function.

        Parameters
        ----------
        assets : List[str]
            List of files to load.
        open_f : Callable[[str], io.IOBase]
            Function to open the files.
            The function takes a file path and returns an io.IOBase object.
        """
        split_files = []
        for path in assets:
            splitMatch = reSplit.match(path)
            if splitMatch:
                basepath, basename = splitMatch.groups()

                if basepath in split_files:
                    continue

                split_files.append(basepath)
                data = []
                for i in range(0, 999):
                    item = f"{basepath}.split{i}"
                    if item in assets:
                        with open_f(item) as f:
                            data.append(f.read())
                    elif i:
                        break
                data = b"".join(data)
                path = basepath
            else:
                data = open_f(path).read()
            self.load_file(data, name=path)

    def find_file(self, name: str, is_dependency: bool = True) -> Union[File, None]:
        """
        Finds a file in the environment.

        Parameters
        ----------
        name : str
            The name of the file.
        is_dependency : bool
            Whether the file is a dependency.

        Returns
        -------
        File | None
            The file if it was found, otherwise None.
        """
        simple_name = simplify_name(name)
        cab = self.get_cab(simple_name)
        if cab:
            return cab

        if len(self.local_files) == 0 and self.path:
            for root, _, files in self.fs.walk(self.path):
                for name in files:
                    self.local_files.append(self.fs.sep.join([root, name]))

        if name in self.local_files:
            fp = name
        elif simple_name in self.local_files_simple:
            fp = self.local_files[self.local_files_simple.index(simple_name)]
        else:
            raise FileNotFoundError(f"File {name} not found in {self.path}")

        f = self.load_file(fp, name=name, is_dependency=is_dependency)
        return f


def simplify_name(name: str) -> str:
    """Simplifies a name by:
    - removing the extension
    - removing the path
    - converting to lowercase
    """
    return ntpath.basename(name).lower()