AvaPackerModule.cpp

// Copyright Epic Games, Inc. All Rights Reserved.

#include "IAvaPackerModule.h"

#include "AssetRegistryModule.h"
#include "IAssetTools.h"
#include "AssetToolsModule.h"

#include "GenericPlatform/GenericPlatformFile.h"
#include "Algo/Accumulate.h"
#include "Misc/KeyChainUtilities.h"


DEFINE_LOG_CATEGORY(AvaPackerLog);

PRAGMA_DISABLE_OPTIMIZATION

class FAvaPackerModule
	: public IAvaPackerModule
{
public:
	virtual void StartupModule() override
	{
		RegisterConsoleCommands();
	}

	virtual void ShutdownModule() override
	{
		UnregisterConsoleCommands();
	}

private:
	/** Register console commands. */
	void RegisterConsoleCommands()
	{
		ConsoleCommands.Add(MakeUnique<FAutoConsoleCommand>(
			TEXT("AvaPacker.Pack"),
			TEXT("Pack"),
			FConsoleCommandWithArgsDelegate::CreateRaw(this, &FAvaPackerModule::Pack)
			));

		ConsoleCommands.Add(MakeUnique<FAutoConsoleCommand>(
			TEXT("AvaPacker.Unpack"),
			TEXT("Unpack"),
			FConsoleCommandWithArgsDelegate::CreateRaw(this, &FAvaPackerModule::Unpack)
			));
	}

	/** Unregister console commands. */
	void UnregisterConsoleCommands()
	{
		for (TUniquePtr<FAutoConsoleCommand>& Command : ConsoleCommands)
		{
			Command.Reset();
		}
	}

	/** Gets the dependencies of the specified package recursively */
	void RecursiveGetDependencies(const FName& PackageName, TSet<FName>& AllDependencies, TSet<FString>& ExternalObjectsPaths) const
	{
		FAssetRegistryModule& AssetRegistryModule = FModuleManager::Get().LoadModuleChecked<FAssetRegistryModule>(TEXT("AssetRegistry"));
		TArray<FName> Dependencies;
		AssetRegistryModule.Get().GetDependencies(PackageName, Dependencies);
	
		for ( auto DependsIt = Dependencies.CreateConstIterator(); DependsIt; ++DependsIt )
		{
			FString DependencyName = (*DependsIt).ToString();

			const bool bIsEnginePackage = DependencyName.StartsWith(TEXT("/Engine"));
			const bool bIsScriptPackage = DependencyName.StartsWith(TEXT("/Script"));
			if ( !bIsEnginePackage && !bIsScriptPackage )
			{
				if (!AllDependencies.Contains(*DependsIt))
				{
					AllDependencies.Add(*DependsIt);
					RecursiveGetDependencies(*DependsIt, AllDependencies, ExternalObjectsPaths);
				}
			}
		}
	}
	

	// C:\Engines\ue5-main\Engine\Source\Developer\AssetTools\Private\AssetTools.cpp
	// C:\Engines\ue5-main\Engine\Source\Developer\PakFileUtilities\Private\PakFileUtilities.cpp
	void Pack(const TArray<FString>& Args)
	{
		// Load all the assets in the selected paths
		FAssetToolsModule& AssetToolsModule = FModuleManager::Get().LoadModuleChecked<FAssetToolsModule>("AssetTools");

		if (Args.Num() > 1)
		{
			FString Package = Args[0];
			const FString PackFileName = Args[1];

			TSet<FName> AllPackageNamesToPak;
			TSet<FString> ExternalObjectsPaths;
			AllPackageNamesToPak.Add(*Package);
			RecursiveGetDependencies(*Package, AllPackageNamesToPak, ExternalObjectsPaths);

			// Create Pak
			TUniquePtr<FArchive> PakFileHandle(IFileManager::Get().CreateFileWriter(*PackFileName));
			if (!PakFileHandle)
			{
				UE_LOG(AvaPackerLog, Error, TEXT("Unable to create pak file \"%s\"."), *PackFileName);
				return;
			}

			int32 AllPackageNamesToPakNum = AllPackageNamesToPak.Num();
			// Write num files in pack
			*PakFileHandle << AllPackageNamesToPakNum;
			
			for (const FName& PackageName : AllPackageNamesToPak)
			{
				FString SrcFilename;

				FString PackageString = PackageName.ToString();
				if (FPackageName::DoesPackageExist(PackageString, &SrcFilename))
				{
					TUniquePtr<FArchive> FileHandle(IFileManager::Get().CreateFileReader(*SrcFilename));
					bool bFileExists = FileHandle.IsValid();
					if (bFileExists)
					{
						int64 FileSize = FileHandle->TotalSize();
						uint8* FileBuffer = (uint8*)FMemory::Malloc(FileSize);
						FileHandle->Serialize(FileBuffer, FileSize);

						// Write Package path
						FString FinalAssetName = PackageString + TEXT(".uasset");
						*PakFileHandle << FinalAssetName;

						// Write size of the buffer
						*PakFileHandle << FileSize;

						// Write to pack
						PakFileHandle->Serialize(FileBuffer, FileSize);

						FMemory::Free(FileBuffer);
					}
				}
				else
				{
					check(0);
				}
			}
			
			
			PakFileHandle->Close();
			PakFileHandle.Reset();
		}
	}

	// Just for testing, by default it should go to /Content folder
	void Unpack(const TArray<FString>& Args)
	{
		if (Args.Num() > 1)
		{
			FString PackPath = Args[0];
			FString RootFolderPath = Args[1];

			TUniquePtr<FArchive> PackHandle(IFileManager::Get().CreateFileReader(*PackPath));
			bool bFileExists = PackHandle.IsValid();
			if (bFileExists)
			{
				int32 AllPackageNamesToPakNum = 0;
				*PackHandle << AllPackageNamesToPakNum;

				for (int32 Index = 0; Index < AllPackageNamesToPakNum; ++Index)
				{
					FString PackagePath;
					*PackHandle << PackagePath;
					
					int64 FileSize = 0;
					*PackHandle << FileSize;

					uint8* FileBuffer = (uint8*)FMemory::Malloc(FileSize);
					PackHandle->Serialize(FileBuffer, FileSize);

					FString FinalAssetPath = FPaths::Combine(RootFolderPath, PackagePath);
					
					TUniquePtr<FArchive> AssetHandle(IFileManager::Get().CreateFileWriter(*FinalAssetPath));

					if (AssetHandle)
					{
						// Write to asset
						AssetHandle->Serialize(FileBuffer, FileSize);

						AssetHandle->Close();
						AssetHandle.Reset();
					}

					FMemory::Free(FileBuffer);
				}

			}
		}
	}

private:
	/** Console commands handles. */
	TArray<TUniquePtr<FAutoConsoleCommand>> ConsoleCommands;
};

//RecursiveGetDependencies

PRAGMA_ENABLE_OPTIMIZATION

IMPLEMENT_MODULE(FAvaPackerModule, AvaPacker);