From e63e3be3b7e0cc89724e7aaee064d4a227fbf671 Mon Sep 17 00:00:00 2001 From: Timothy Baldridge Date: Fri, 26 Jul 2019 14:59:14 -0600 Subject: [PATCH] working on BSA compression support --- BSA.Tools/Archive.cs | 409 ----------------------------------- Wabbajack.Common/BSDiff.cs | 34 ++- Wabbajack.Common/Consts.cs | 2 + Wabbajack.Common/Data.cs | 21 +- Wabbajack.Common/Utils.cs | 8 + Wabbajack/App.config | 10 +- Wabbajack/Compiler.cs | 153 ++++++++++++- Wabbajack/Installer.cs | 52 +++-- Wabbajack/MainWindow.xaml.cs | 9 +- 9 files changed, 256 insertions(+), 442 deletions(-) delete mode 100644 BSA.Tools/Archive.cs diff --git a/BSA.Tools/Archive.cs b/BSA.Tools/Archive.cs deleted file mode 100644 index 6460a1a4..00000000 --- a/BSA.Tools/Archive.cs +++ /dev/null @@ -1,409 +0,0 @@ -using System; -using System.Collections.Generic; -using System.IO; -using System.Runtime.InteropServices; -using System.Text; -using static BSA.Tools.libbsarch; - -namespace BSA.Tools -{ - // Represents a BSA archive on disk (in READ mode) - public class Archive : IDisposable - { - protected unsafe libbsarch.bsa_archive_t* _archive; - - public UInt32 Version - { - get - { - lock(this) { - unsafe - { - return libbsarch.bsa_version_get(_archive); - } - } - } - } - - public bsa_archive_type_t Type - { - get - { - lock(this) - { - unsafe - { - return libbsarch.bsa_archive_type_get(_archive); - } - } - } - } - - public UInt32 FileCount - { - get - { - lock (this) - { - unsafe - { - return libbsarch.bsa_file_count_get(_archive); - } - } - } - } - - public UInt32 ArchiveFlags - { - get - { - lock (this) - { - unsafe - { - return libbsarch.bsa_archive_flags_get(_archive); - } - } - } - set - { - lock (this) - { - unsafe - { - libbsarch.bsa_archive_flags_set(_archive, value); - } - } - } - } - - public UInt32 FileFlags - { - get - { - lock (this) - { - unsafe - { - return libbsarch.bsa_file_flags_get(_archive); - } - } - } - set - { - lock (this) - { - unsafe - { - libbsarch.bsa_file_flags_set(_archive, value); - } - } - } - } - - public bool Compress - { - get - { - lock (this) - { - unsafe - { - return libbsarch.bsa_compress_get(_archive); - } - } - } - set - { - lock (this) - { - unsafe - { - libbsarch.bsa_compress_set(_archive, value); - } - } - } - } - - public bool ShareData - { - get - { - lock (this) - { - unsafe - { - return libbsarch.bsa_share_data_get(_archive); - } - } - } - set - { - lock (this) - { - unsafe - { - libbsarch.bsa_share_data_set(_archive, value); - } - } - } - } - - - public void Save() - { - lock (this) - { - unsafe - { - check_err(libbsarch.bsa_save(_archive)); - } - } - } - - private IEnumerable _entries = null; - public IEnumerable Entries { - get - { - if (_entries != null) - return _entries; - - return GetAndCacheEntries(); - } - - - } - - private IEnumerable GetAndCacheEntries() - { - var entries = new List(); - unsafe - { - foreach (var filename in GetFileNames()) - { - entries.Add(new ArchiveEntry(this, _archive, filename)); - } - } - _entries = entries; - return entries; - } - - public Archive() - { - unsafe - { - _archive = libbsarch.bsa_create(); - } - } - - public void Create(string filename, bsa_archive_type_t type, EntryList entries) - { - unsafe - { - check_err(libbsarch.bsa_create_archive(_archive, filename, type, entries._list)); - } - } - - public Archive(string filename) - { - unsafe - { - _archive = libbsarch.bsa_create(); - check_err(libbsarch.bsa_load_from_file(_archive, filename)); - } - } - - public void AddFile(string filename, byte[] data) - { - lock(this) - { - unsafe - { - var ptr = Marshal.AllocHGlobal(data.Length); - Marshal.Copy(data, 0, ptr, data.Length); - libbsarch.bsa_add_file_from_memory(_archive, filename, (UInt32)data.Length, (byte*)ptr); - Marshal.FreeHGlobal(ptr); - } - } - } - - public void Dispose() - { - unsafe - { - check_err(libbsarch.bsa_free(_archive)); - } - } - - public static void check_err(libbsarch.bsa_result_message_t bsa_result_message_t) - { - if (bsa_result_message_t.code != 0) - { - unsafe - { - int i = 0; - for (i = 0; i < 1024 * 2; i += 2) - if (bsa_result_message_t.text[i] == 0) break; - - var msg = new String((sbyte*)bsa_result_message_t.text, 0, i, Encoding.Unicode); - throw new Exception(msg); - } - } - } - - public IEnumerable GetFileNames() - { - List filenames = new List(); - lock (this) - { - unsafe - { - check_err(libbsarch.bsa_iterate_files(_archive, (archive, filename, file, folder, context) => - { - lock (filenames) - { - filenames.Add(filename); - } - return false; - }, null)); - } - } - return filenames; - } - } - - public class ArchiveEntry - { - private Archive _archive; - private unsafe libbsarch.bsa_archive_t* _archivep; - private string _filename; - - public string Filename { - get - { - return _filename; - } - } - - public unsafe ArchiveEntry(Archive archive, libbsarch.bsa_archive_t* archivep, string filename) - { - _archive = archive; - _archivep = archivep; - _filename = filename; - } - - public FileData GetFileData() - { - unsafe - { - var result = libbsarch.bsa_extract_file_data_by_filename(_archivep, _filename); - Archive.check_err(result.message); - return new FileData(_archive, _archivep, result.buffer); - } - } - - public void ExtractTo(Stream stream) - { - using (var data = GetFileData()) - { - data.WriteTo(stream); - } - } - - public void ExtractTo(string filename) - { - unsafe - { - libbsarch.bsa_extract_file(_archivep, _filename, filename); - } - } - } - - public class FileData : IDisposable - { - private Archive archive; - private unsafe libbsarch.bsa_archive_t* archivep; - private libbsarch.bsa_result_buffer_t result; - - public unsafe FileData(Archive archive, libbsarch.bsa_archive_t* archivep, libbsarch.bsa_result_buffer_t result) - { - this.archive = archive; - this.archivep = archivep; - this.result = result; - } - - public void WriteTo(Stream stream) - { - var memory = ToByteArray(); - stream.Write(memory, 0, (int)result.size); - } - - public byte[] ToByteArray() - { - unsafe - { - byte[] memory = new byte[result.size]; - Marshal.Copy((IntPtr)result.data, memory, 0, (int)result.size); - return memory; - } - } - - public void Dispose() - { - lock(archive) - unsafe - { - Archive.check_err(libbsarch.bsa_file_data_free(archivep, result)); - } - } - } - - public class EntryList : IDisposable - { - public unsafe bsa_entry_list_t* _list; - - public EntryList() - { - unsafe - { - _list = libbsarch.bsa_entry_list_create(); - } - } - - public UInt32 Count - { - get - { - lock (this) - { - unsafe - { - return libbsarch.bsa_entry_list_count(_list); - } - } - } - } - - public void Add(string entry) - { - lock(this) - { - unsafe - { - libbsarch.bsa_entry_list_add(_list, entry); - } - } - } - - public void Dispose() - { - lock (this) - { - unsafe - { - libbsarch.bsa_entry_list_free(_list); - } - } - } - } -} diff --git a/Wabbajack.Common/BSDiff.cs b/Wabbajack.Common/BSDiff.cs index c51eaa19..26c16b52 100644 --- a/Wabbajack.Common/BSDiff.cs +++ b/Wabbajack.Common/BSDiff.cs @@ -31,6 +31,11 @@ IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */ + + /* + * Includes fixes from https://github.com/mendsley/bsdiff/pull/12/files + * translated from C to C# to prevent stack overflows. + */ public class BSDiff { /// @@ -444,7 +449,23 @@ } else { - int x = v[I[start + len / 2] + h]; + int y, z, j, k, x, tmp; + j = start + len / 2; + k = start + len - 1; + x = v[I[j] + h]; + y = v[I[start] + h]; + z = v[I[k] + h]; + if (len > 40) + { /* Big array: Pseudomedian of 9 */ + tmp = len / 8; + x = median3(x, v[I[j - tmp] + h], v[I[j + tmp] + h]); + y = median3(y, v[I[start + tmp] + h], v[I[start + tmp + tmp] + h]); + z = median3(z, v[I[k - tmp] + h], v[I[k - tmp - tmp] + h]); + }; /* Else medium array: Pseudomedian of 3 */ + x = median3(x, y, z); + + //int x = v[I[start + len / 2] + h]; + int jj = 0; int kk = 0; for (int i2 = start; i2 < start + len; i2++) @@ -458,8 +479,8 @@ kk += jj; int i = start; - int j = 0; - int k = 0; + j = 0; + k = 0; while (i < jj) { if (v[I[i] + h] < x) @@ -652,5 +673,12 @@ const long c_fileSignature = 0x3034464649445342L; const int c_headerSize = 32; + + + private static int median3(int a, int b, int c) + { + return (((a) < (b)) ? ((b) < (c) ? (b) : ((a) < (c) ? (c) : (a))) : ((b) > (c) ? (b) : ((a) > (c) ? (c) : (a)))); + } + } } diff --git a/Wabbajack.Common/Consts.cs b/Wabbajack.Common/Consts.cs index dcda85bb..6180e516 100644 --- a/Wabbajack.Common/Consts.cs +++ b/Wabbajack.Common/Consts.cs @@ -12,8 +12,10 @@ namespace Wabbajack.Common { public static string GameFolderFilesDir = "Game Folder Files"; public static string ModPackMagic = "Celebration!, Cheese for Everyone!"; + public static string BSACreationDir = "TEMP_BSA_FILES"; public static HashSet SupportedArchives = new HashSet() { ".zip", ".rar", ".7z", ".7zip" }; + public static HashSet SupportedBSAs = new HashSet() { ".bsa", ".ba2" }; public static String UserAgent { get diff --git a/Wabbajack.Common/Data.cs b/Wabbajack.Common/Data.cs index c9653283..7b8b2785 100644 --- a/Wabbajack.Common/Data.cs +++ b/Wabbajack.Common/Data.cs @@ -20,6 +20,10 @@ namespace Wabbajack.Common _hash = AbsolutePath.FileSHA256(); return _hash; } + set + { + _hash = value; + } } public T EvolveTo() where T : Directive, new() @@ -94,6 +98,17 @@ namespace Wabbajack.Common public string From; } + public class CreateBSA : Directive + { + public string TempID; + public string IsCompressed; + public uint Version; + public Int32 Type; + + public uint FileFlags { get; set; } + public bool Compress { get; set; } + } + public class PatchedFromArchive : FromArchive { /// @@ -113,7 +128,6 @@ namespace Wabbajack.Common /// public string Name; - /// /// Meta INI for the downloaded archive /// public string Meta; @@ -126,6 +140,11 @@ namespace Wabbajack.Common public string FileID; } + public class GoogleDriveMod : Archive + { + public string Id; + } + /// /// URL that can be downloaded directly without any additional options /// diff --git a/Wabbajack.Common/Utils.cs b/Wabbajack.Common/Utils.cs index fa080bbb..c81fa061 100644 --- a/Wabbajack.Common/Utils.cs +++ b/Wabbajack.Common/Utils.cs @@ -34,6 +34,12 @@ namespace Wabbajack.Common } + public static string SHA256(this byte[] data) + { + return new SHA256Managed().ComputeHash(data).ToBase64(); + + } + /// /// Returns a Base64 encoding of these bytes /// @@ -167,6 +173,8 @@ namespace Wabbajack.Common return tasks.Select(t => { t.Wait(); + if (t.IsFaulted) + throw t.Exception; return t.Result; }).ToList(); } diff --git a/Wabbajack/App.config b/Wabbajack/App.config index 56efbc7b..892c7c30 100644 --- a/Wabbajack/App.config +++ b/Wabbajack/App.config @@ -1,6 +1,14 @@ - + + + + + + + + + \ No newline at end of file diff --git a/Wabbajack/Compiler.cs b/Wabbajack/Compiler.cs index 61a58784..f748b9f3 100644 --- a/Wabbajack/Compiler.cs +++ b/Wabbajack/Compiler.cs @@ -1,6 +1,8 @@ -using Newtonsoft.Json; +using BSA.Tools; +using Newtonsoft.Json; using SevenZipExtractor; using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.IO; using System.Linq; @@ -49,6 +51,7 @@ namespace Wabbajack public List SelectedArchives { get; private set; } public List AllFiles { get; private set; } public ModList ModList { get; private set; } + public ConcurrentBag ExtraFiles { get; private set; } public List IndexedArchives; @@ -159,15 +162,23 @@ namespace Wabbajack Info("Found {0} files to build into mod list", AllFiles.Count); + ExtraFiles = new ConcurrentBag(); + var stack = MakeStack(); Info("Running Compilation Stack"); var results = AllFiles.PMap(f => RunStack(stack, f)).ToList(); + // Add the extra files that were generated by the stack + Info($"Adding {ExtraFiles.Count} that were generated by the stack"); + results = results.Concat(ExtraFiles).ToList(); + var nomatch = results.OfType(); Info("No match for {0} files", nomatch.Count()); foreach (var file in nomatch) Info(" {0}", file.To); + if (nomatch.Count() > 0) + Error("Exiting due to no way to compile these files"); InstallDirectives = results.Where(i => !(i is IgnoredDirectly)).ToList(); @@ -196,6 +207,7 @@ namespace Wabbajack IndexedArchives = null; InstallDirectives = null; SelectedArchives = null; + ExtraFiles = null; } @@ -224,7 +236,7 @@ namespace Wabbajack var archive = IndexedArchives.First(a => a.Hash == archive_sha); var paths = group.Select(g => g.From).ToHashSet(); var streams = new Dictionary(); - Status("Etracting Patch Files from {0}", archive.Name); + Status($"Extracting {paths.Count} patch files from {archive.Name}"); // First we fetch the source files from the input archive using (var a = new ArchiveFile(archive.AbsolutePath)) { @@ -246,11 +258,10 @@ namespace Wabbajack Info("Patching {0}", entry.To); var ss = extracted[entry.From]; using (var origin = new MemoryStream(ss)) - using (var dest = File.OpenRead(absolute_paths[entry.To])) using (var output = new MemoryStream()) { var a = origin.ReadAll(); - var b = dest.ReadAll(); + var b = LoadDataForTo(entry.To, absolute_paths); BSDiff.Create(a, b, output); entry.Patch = output.ToArray().ToBase64(); } @@ -258,6 +269,30 @@ namespace Wabbajack } + private byte[] LoadDataForTo(string to, Dictionary absolute_paths) + { + if (absolute_paths.TryGetValue(to, out var absolute)) + return File.ReadAllBytes(absolute); + + if (to.StartsWith(Consts.BSACreationDir)) + { + var bsa_id = to.Split('\\')[1]; + var bsa = InstallDirectives.OfType().First(b => b.TempID == bsa_id); + + using (var a = new BSAFile(Path.Combine(MO2Folder, bsa.To))) + { + var file = a.Entries.First(e => e.Filename == Path.Combine(to.Split('\\').Skip(2).ToArray())); + using (var data = file.GetFileData()) + { + return data.ToByteArray(); + } + } + + } + Error($"Couldn't load data for {to}"); + return null; + } + private void GatherArchives() { var archives = IndexedArchives.GroupBy(a => a.Hash).ToDictionary(k => k.Key, k => k.First()); @@ -291,6 +326,15 @@ namespace Wabbajack ModID = general.modID }; } + else if (general.directURL != null && general.directURL.StartsWith("https://drive.google.com")) + { + var regex = new Regex("((?<=id=)[a-zA-Z0-9_-]*)|(?<=\\/file\\/d\\/)[a-zA-Z0-9_-]*"); + var match = regex.Match(general.directURL); + result = new GoogleDriveMod() + { + Id = match.ToString() + }; + } else if (general.directURL != null && general.directURL.StartsWith("https://www.dropbox.com/")) { var uri = new UriBuilder((string)general.directURL); @@ -363,6 +407,7 @@ namespace Wabbajack IgnoreStartsWith("logs\\"), IgnoreStartsWith("downloads\\"), IgnoreStartsWith("webcache\\"), + IgnoreStartsWith("overwrite\\"), IgnoreEndsWith(".pyc"), IgnoreOtherProfiles(), IgnoreDisabledMods(), @@ -377,12 +422,108 @@ namespace Wabbajack DirectMatch(), IncludePatches(), + DeconstructBSAs(), + // If we have no match at this point for a game folder file, skip them, we can't do anything about them IgnoreGameFiles(), + + // There are some types of files that will error the compilation, because tehy're created on-the-fly via tools + // so if we don't have a match by this point, just drop them. + IgnoreEndsWith(".ini"), + IgnoreEndsWith(".html"), + IgnoreEndsWith(".txt"), + // Don't know why, but this seems to get copied around a bit + IgnoreEndsWith("HavokBehaviorPostProcess.exe"), DropAll() }; } + + /// + /// This function will search for a way to create a BSA in the installed mod list by assembling it from files + /// found in archives. To do this we hash all the files in side the BSA then try to find matches and patches for + /// all of the files. + /// + /// + private Func DeconstructBSAs() + { + var microstack = new List>() + { + DirectMatch(), + IncludePatches(), + DropAll() + }; + + return source => + { + if (!Consts.SupportedBSAs.Contains(Path.GetExtension(source.Path))) return null; + + var hashed = HashBSA(source.AbsolutePath); + + var source_files = hashed.Select(e => new RawSourceFile() { + Hash = e.Item2, + Path = e.Item1, + AbsolutePath = e.Item1 + }); + + + var matches = source_files.Select(e => RunStack(microstack, e)); + + var id = Guid.NewGuid().ToString(); + + foreach (var match in matches) + { + if (match is IgnoredDirectly) + { + Error($"File required for BSA creation doesn't exist: {match.To}"); + } + match.To = Path.Combine(Consts.BSACreationDir, id, match.To); + ExtraFiles.Add(match); + }; + + CreateBSA directive; + using (var bsa = new BSAFile(source.AbsolutePath)) + { + directive = new CreateBSA() + { + To = source.Path, + TempID = id, + Version = bsa.Version, + Type = (int)bsa.Type, + FileFlags = bsa.FileFlags, + Compress = bsa.Compress + }; + }; + + return directive; + + }; + } + + /// + /// Given a BSA on disk, index it and return a dictionary of SHA256 -> filename + /// + /// + /// + private List<(string, string)> HashBSA(string absolutePath) + { + Status($"Hashing BSA: {absolutePath}"); + var results = new List<(string, string)>(); + using (var a = new BSAFile(absolutePath)) + { + foreach (var entry in a.Entries) + { + Status($"Hashing BSA: {absolutePath} - {entry.Filename}"); + + using (var data = entry.GetFileData()) + { + results.Add((entry.Filename, data.ToByteArray().SHA256())); + }; + } + } + return results; + } + private Func IgnoreDisabledMods() { var disabled_mods = File.ReadAllLines(Path.Combine(MO2ProfileDir, "modlist.txt")) @@ -406,12 +547,12 @@ namespace Wabbajack var indexed = (from archive in IndexedArchives from entry in archive.Entries select new { archive = archive, entry = entry }) - .GroupBy(e => Path.GetFileName(e.entry.Path)) + .GroupBy(e => Path.GetFileName(e.entry.Path).ToLower()) .ToDictionary(e => e.Key); return source => { - if (indexed.TryGetValue(Path.GetFileName(source.Path), out var value)) + if (indexed.TryGetValue(Path.GetFileName(source.Path.ToLower()), out var value)) { var found = value.First(); diff --git a/Wabbajack/Installer.cs b/Wabbajack/Installer.cs index aa4c30e7..e45bf5b3 100644 --- a/Wabbajack/Installer.cs +++ b/Wabbajack/Installer.cs @@ -31,6 +31,7 @@ namespace Wabbajack public ModList ModList { get; } public Action Log_Fn { get; } public Dictionary HashedArchives { get; private set; } + public string NexusAPIKey { get; private set; } public void Info(string msg, params object[] args) @@ -199,26 +200,37 @@ namespace Wabbajack { missing.PMap(archive => { - if (archive is NexusMod) - { - var url = NexusAPI.GetNexusDownloadLink(archive as NexusMod, NexusAPIKey); - DownloadURLDirect(archive, url); - } - else if (archive is MODDBArchive) - { - DownloadModDBArchive(archive, (archive as MODDBArchive).URL); - } - else if (archive is DirectURLArchive) - { - DownloadURLDirect(archive, (archive as DirectURLArchive).URL); - } - else - { + switch (archive) { + case NexusMod a: + var url = NexusAPI.GetNexusDownloadLink(a as NexusMod, NexusAPIKey); + DownloadURLDirect(archive, url); + break; + case GoogleDriveMod a: + DownloadGoogleDriveArchive(a); + break; + case MODDBArchive a: + DownloadModDBArchive(archive, (archive as MODDBArchive).URL); + break; + case DirectURLArchive a: + DownloadURLDirect(archive, (archive as DirectURLArchive).URL); + break; + default: + break; - } + } }); } + private void DownloadGoogleDriveArchive(GoogleDriveMod a) + { + var initial_url = $"https://drive.google.com/uc?id={a.Id}&export=download"; + var client = new HttpClient(); + var result = client.GetStringSync(initial_url); + var regex = new Regex("(?<=/uc\\?export=download&confirm=).*(?=;id=)"); + var confirm = regex.Match(result); + DownloadURLDirect(a, $"https://drive.google.com/uc?export=download&confirm={confirm}&id={a.Id}", client); + } + private void DownloadModDBArchive(Archive archive, string url) { var client = new HttpClient(); @@ -228,10 +240,12 @@ namespace Wabbajack DownloadURLDirect(archive, match.Value); } - private void DownloadURLDirect(Archive archive, string url) + private void DownloadURLDirect(Archive archive, string url, HttpClient client = null) { - HttpClient client = new HttpClient(); - client.DefaultRequestHeaders.Add("User-Agent", Consts.UserAgent); + if (client == null) { + client = new HttpClient(); + client.DefaultRequestHeaders.Add("User-Agent", Consts.UserAgent); + } long total_read = 0; int buffer_size = 1024 * 32; diff --git a/Wabbajack/MainWindow.xaml.cs b/Wabbajack/MainWindow.xaml.cs index 6adc9e53..62b6cd47 100644 --- a/Wabbajack/MainWindow.xaml.cs +++ b/Wabbajack/MainWindow.xaml.cs @@ -1,5 +1,6 @@ using System; using System.Collections.Generic; +using System.IO; using System.Linq; using System.Text; using System.Threading; @@ -25,6 +26,7 @@ namespace Wabbajack public MainWindow() { InitializeComponent(); + var context = new AppState(Dispatcher, "Building"); WorkQueue.Init((id, msg, progress) => context.SetProgress(id, msg, progress)); @@ -32,12 +34,13 @@ namespace Wabbajack new Thread(() => { compiler.LoadArchives(); - compiler.MO2Profile = "Basic Graphics and Fixes"; + compiler.MO2Profile = "DEV"; //"Basic Graphics and Fixes"; compiler.Compile(); compiler.ModList.ToJSON("C:\\tmp\\modpack.json"); - - var installer = new Installer(compiler.ModList, "c:\\tmp\\install\\", msg => context.LogMsg(msg)); + var modlist = compiler.ModList; + compiler = null; + var installer = new Installer(modlist, "c:\\tmp\\install\\", msg => context.LogMsg(msg)); installer.Install(); }).Start();