From a00128db35ddc928f9b19240edddf919f6d48b46 Mon Sep 17 00:00:00 2001 From: Timothy Baldridge Date: Tue, 21 Jul 2020 16:25:34 -0600 Subject: [PATCH] Check file signatures before attempting to extract. Only use 7zip for RAR, 7zip and Zip (not TAR) --- Compression.BSA/BSADispatch.cs | 25 ++++++-------- Compression.BSA/BSAReader.cs | 2 +- Wabbajack.Lib/Wabbajack.Lib.csproj | 2 +- Wabbajack.Test/EndToEndTests.cs | 3 +- Wabbajack.VirtualFileSystem/FileExtractor.cs | 34 ++++++++++++++++---- 5 files changed, 40 insertions(+), 26 deletions(-) diff --git a/Compression.BSA/BSADispatch.cs b/Compression.BSA/BSADispatch.cs index fec3073c..df8140fd 100644 --- a/Compression.BSA/BSADispatch.cs +++ b/Compression.BSA/BSADispatch.cs @@ -3,6 +3,7 @@ using System.IO; using System.Text; using System.Threading.Tasks; using Wabbajack.Common; +using Wabbajack.Common.FileSignatures; namespace Compression.BSA { @@ -10,27 +11,19 @@ namespace Compression.BSA { public static async ValueTask OpenRead(AbsolutePath filename) { - var fourcc = ""; - using (var file = await filename.OpenRead()) + return await BSASignatures.MatchesAsync(filename) switch { - fourcc = Encoding.ASCII.GetString(new BinaryReader(file).ReadBytes(4)); - } - - if (fourcc == TES3Reader.TES3_MAGIC) - return await TES3Reader.Load(filename); - if (fourcc == "BSA\0") - return await BSAReader.LoadWithRetry(filename); - if (fourcc == "BTDX") - return await BA2Reader.Load(filename); - throw new InvalidDataException("Filename is not a .bsa or .ba2, magic " + fourcc); + Definitions.FileType.TES3 => await TES3Reader.Load(filename), + Definitions.FileType.BSA => await BSAReader.LoadAsync(filename), + Definitions.FileType.BA2 => await BA2Reader.Load(filename), + _ => throw new InvalidDataException("Filename is not a .bsa or .ba2") + }; } - private static HashSet MagicStrings = new HashSet {TES3Reader.TES3_MAGIC, "BSA\0", "BTDX"}; + private static SignatureChecker BSASignatures = new SignatureChecker(Definitions.FileType.BSA, Definitions.FileType.BA2, Definitions.FileType.TES3); public static async ValueTask MightBeBSA(AbsolutePath filename) { - using var file = await filename.OpenRead(); - var fourcc = Encoding.ASCII.GetString(new BinaryReader(file).ReadBytes(4)); - return MagicStrings.Contains(fourcc); + return await BSASignatures.MatchesAsync(filename) != null; } } } diff --git a/Compression.BSA/BSAReader.cs b/Compression.BSA/BSAReader.cs index 452c43b4..f580555a 100644 --- a/Compression.BSA/BSAReader.cs +++ b/Compression.BSA/BSAReader.cs @@ -77,7 +77,7 @@ namespace Compression.BSA } } - public static async ValueTask LoadWithRetry(AbsolutePath filename) + public static async ValueTask LoadAsync(AbsolutePath filename) { using var stream = await filename.OpenRead(); using var br = new BinaryReader(stream); diff --git a/Wabbajack.Lib/Wabbajack.Lib.csproj b/Wabbajack.Lib/Wabbajack.Lib.csproj index 1e80fd0a..8a303c17 100644 --- a/Wabbajack.Lib/Wabbajack.Lib.csproj +++ b/Wabbajack.Lib/Wabbajack.Lib.csproj @@ -68,7 +68,7 @@ 1.0.0 - 5.1.1 + 5.1.2 diff --git a/Wabbajack.Test/EndToEndTests.cs b/Wabbajack.Test/EndToEndTests.cs index 198da817..879a7b3d 100644 --- a/Wabbajack.Test/EndToEndTests.cs +++ b/Wabbajack.Test/EndToEndTests.cs @@ -61,7 +61,8 @@ namespace Wabbajack.Test DownloadAndInstall(Game.SkyrimSpecialEdition, 12604, "SkyUI"), DownloadAndInstall(Game.Fallout4, 11925, "Anti-Tank Rifle"), DownloadAndInstall(Game.SkyrimSpecialEdition, 4783, "Frost Armor UNP"), - DownloadAndInstall(Game.SkyrimSpecialEdition, 32359, "Frost Armor HDT")); + DownloadAndInstall(Game.SkyrimSpecialEdition, 32359, "Frost Armor HDT"), + DownloadAndInstall(Game.SkyrimSpecialEdition, 31667, "Nemesis")); // We're going to fully patch this mod from another source. await modfiles[3].Download.DeleteAsync(); diff --git a/Wabbajack.VirtualFileSystem/FileExtractor.cs b/Wabbajack.VirtualFileSystem/FileExtractor.cs index 1fd4ffae..cba79bf1 100644 --- a/Wabbajack.VirtualFileSystem/FileExtractor.cs +++ b/Wabbajack.VirtualFileSystem/FileExtractor.cs @@ -10,6 +10,7 @@ using OMODFramework; using Wabbajack.Common.StatusFeed; using Wabbajack.Common.StatusFeed.Errors; using Wabbajack.Common; +using Wabbajack.Common.FileSignatures; using Utils = Wabbajack.Common.Utils; @@ -17,19 +18,38 @@ namespace Wabbajack.VirtualFileSystem { public class FileExtractor { + + private static SignatureChecker archiveSigs = new SignatureChecker(Definitions.FileType.TES3, + Definitions.FileType.BSA, + Definitions.FileType.BA2, + Definitions.FileType.ZIP, + Definitions.FileType.EXE, + Definitions.FileType.RAR, + Definitions.FileType._7Z); public static async Task ExtractAll(WorkQueue queue, AbsolutePath source, IEnumerable OnlyFiles = null) { try { - if (await BSADispatch.MightBeBSA(source)) - return await ExtractAllWithBSA(queue, source); - else if (source.Extension == Consts.OMOD) + var sig = await archiveSigs.MatchesAsync(source); + + if (source.Extension == Consts.OMOD) return await ExtractAllWithOMOD(source); - else if (source.Extension == Consts.EXE) - return await ExtractAllExe(source); - else - return await ExtractAllWith7Zip(source, OnlyFiles); + + switch (sig) + { + case Definitions.FileType.BSA: + case Definitions.FileType.TES3: + case Definitions.FileType.BA2: + return await ExtractAllWithBSA(queue, source); + case Definitions.FileType.EXE: + return await ExtractAllExe(source); + case Definitions.FileType._7Z: + case Definitions.FileType.ZIP: + case Definitions.FileType.RAR: + return await ExtractAllWith7Zip(source, OnlyFiles); + } + throw new Exception("Invalid archive format"); } catch (Exception ex) {