import io

import ntpath

import os

import re

from typing import TYPE_CHECKING, BinaryIO, Callable, Dict, List, Optional, Union, cast

from zipfile import ZipFile

from fsspec import AbstractFileSystem

from fsspec.implementations.local import LocalFileSystem

from .enums import FileType

from .files import BundleFile, File, ObjectReader, SerializedFile, WebFile

from .helpers.ContainerHelper import ContainerHelper

from .helpers.ImportHelper import (

FileSourceType,

check_file_type,

find_sensitive_path,

parse_file,

)

from .streams import EndianBinaryReader

if TYPE_CHECKING:

from UnityPy.helpers.TypeTreeGenerator import TypeTreeGenerator

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

class Environment:

files: Dict[str, Union[SerializedFile, BundleFile, WebFile, EndianBinaryReader]]

cabs: Dict[str, Union[SerializedFile, EndianBinaryReader]]

path: str

local_files: List[str]

local_files_simple: List[str]

typetree_generator: Optional["TypeTreeGenerator"] = None

_container_index_built: bool = False

def __init__(self, *args: FileSourceType, fs: Optional[AbstractFileSystem] = None, path: Optional[str] = None):

self.files = {}

self.cabs = {}

self.fs = fs or LocalFileSystem()

self.local_files = []

self.local_files_simple = []

self._container_index_built = False

if path is None:

# if no path is given, use the current working directory

if isinstance(self.fs, LocalFileSystem):

self.path = os.getcwd()

else:

self.path = ""

else:

self.path = path

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) or ntpath.curdir

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.load_file(file=arg)

if len(self.files) == 1:

self.file = list(self.files.values())[0]

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[str]):

"""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_split_file(self, basename: str) -> bytes:

file: List[bytes] = []

for i in range(0, 999):

item = f"{basename}.split{i}"

if self.fs.exists(item):

with self.fs.open(item, "rb") as f:

file.append(f.read()) # type: ignore

elif i:

break

return b"".join(file)

def load_file(

self,

file: FileSourceType,

parent: Optional[Union["Environment", File]] = None,

name: Optional[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()

assert isinstance(basepath, str)

name = basepath

file = self._load_split_file(basepath)

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)

# for dependency loading of split files

if os.path.exists(f"{file}.split0"):

file = self._load_split_file(file)

# Unity paths are case insensitive,

# so we need to find "Resources/Foo.asset" when the record says "resources/foo.asset"

elif not os.path.exists(file):

file_path = find_sensitive_path(self.path, file)

if file_path:

file = file_path

else:

return None

# raise FileNotFoundError(f"File {file} not found in {self.path}")

if isinstance(file, str):

file = self.fs.open(file, "rb")

typ, reader = check_file_type(file)

stream_name = (

name

if name

else getattr(

file,

"name",

str(file.__hash__()) if hasattr(file, "__hash__") else "", # type: ignore

)

)

if typ == FileType.ZIP:

f = self.load_zip_file(file)

else:

f = 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

return f

def load_zip_file(self, value):

if isinstance(value, str) and self.fs.exists(value):

buffer = cast(io.BufferedReader, self.fs.open(value, "rb"))

elif isinstance(value, (bytes, bytearray, memoryview)):

buffer = io.BytesIO(value)

elif isinstance(value, (io.BufferedReader, io.BufferedIOBase)):

buffer = value

else:

raise TypeError("Unsupported type for loading zip file")

z = ZipFile(buffer)

self.load_assets(z.namelist(), lambda x: z.open(x, "r")) # type: ignore

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 sub_item in item.files.values():

ret.extend(search(sub_item))

return ret

return ret

return search(self)

def _build_container_index(self) -> None:

if self._container_index_built:

return

self._container_index_built = True

for f in self.cabs.values():

if isinstance(f, SerializedFile):

f.container.parse_preload_table()

@property

def container(self) -> ContainerHelper:

"""Returns a dictionary of all objects in the Environment."""

self._build_container_index()

container = []

for f in self.cabs.values():

if isinstance(f, SerializedFile) and not f.is_dependency:

container.extend(f.container.container)

return ContainerHelper(container)

@property

def assets(self) -> list:

"""

Lists all assets / SerializedFiles within this environment.

"""

def gen_all_asset_files(file, ret: Optional[list] = None):

if ret is None:

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: Union[SerializedFile, EndianBinaryReader]) -> 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

self._container_index_built = False

def get_cab(self, name: str) -> Union[SerializedFile, EndianBinaryReader, None]:

"""

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], BinaryIO]):

"""

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 = self._load_split_file(basepath)

path = basepath

else:

data = open_f(path)

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

fp = self.fs.sep.join([self.path, name])

if self.fs.exists(fp):

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

if len(self.local_files) == 0 and self.path:

for root, _, files in self.fs.walk(self.path):

for f in files:

self.local_files.append(self.fs.sep.join([root, f]))

self.local_files_simple.append(self.fs.sep.join([root, simplify_name(f)]))

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:

fp = next((f for f in self.local_files if f.endswith(name)), None)

if not fp:

fp = next(

(f for f in self.local_files_simple if f.endswith(simple_name)),

None,

)

if not fp:

raise FileNotFoundError(f"File {name} not found in {self.path}")

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

def simplify_name(name: str) -> str:

"""Simplifies a name by:

- removing the extension

- removing the path

- converting to lowercase

"""

return ntpath.basename(name).lower()