Merge pull request #973 from wabbajack-tools/check-file-sigs

Check file signatures before attempting to extract. Only use 7zip for…
This commit is contained in:
Timothy Baldridge 2020-07-21 19:45:52 -07:00 committed by GitHub
commit fefae9cad2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 40 additions and 26 deletions

View File

@ -3,6 +3,7 @@ using System.IO;
using System.Text; using System.Text;
using System.Threading.Tasks; using System.Threading.Tasks;
using Wabbajack.Common; using Wabbajack.Common;
using Wabbajack.Common.FileSignatures;
namespace Compression.BSA namespace Compression.BSA
{ {
@ -10,27 +11,19 @@ namespace Compression.BSA
{ {
public static async ValueTask<IBSAReader> OpenRead(AbsolutePath filename) public static async ValueTask<IBSAReader> OpenRead(AbsolutePath filename)
{ {
var fourcc = ""; return await BSASignatures.MatchesAsync(filename) switch
using (var file = await filename.OpenRead())
{ {
fourcc = Encoding.ASCII.GetString(new BinaryReader(file).ReadBytes(4)); Definitions.FileType.TES3 => await TES3Reader.Load(filename),
} Definitions.FileType.BSA => await BSAReader.LoadAsync(filename),
Definitions.FileType.BA2 => await BA2Reader.Load(filename),
if (fourcc == TES3Reader.TES3_MAGIC) _ => throw new InvalidDataException("Filename is not a .bsa or .ba2")
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);
} }
private static HashSet<string> MagicStrings = new HashSet<string> {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<bool> MightBeBSA(AbsolutePath filename) public static async ValueTask<bool> MightBeBSA(AbsolutePath filename)
{ {
using var file = await filename.OpenRead(); return await BSASignatures.MatchesAsync(filename) != null;
var fourcc = Encoding.ASCII.GetString(new BinaryReader(file).ReadBytes(4));
return MagicStrings.Contains(fourcc);
} }
} }
} }

View File

@ -77,7 +77,7 @@ namespace Compression.BSA
} }
} }
public static async ValueTask<BSAReader> LoadWithRetry(AbsolutePath filename) public static async ValueTask<BSAReader> LoadAsync(AbsolutePath filename)
{ {
using var stream = await filename.OpenRead(); using var stream = await filename.OpenRead();
using var br = new BinaryReader(stream); using var br = new BinaryReader(stream);

View File

@ -68,7 +68,7 @@
<Version>1.0.0</Version> <Version>1.0.0</Version>
</PackageReference> </PackageReference>
<PackageReference Include="YoutubeExplode"> <PackageReference Include="YoutubeExplode">
<Version>5.1.1</Version> <Version>5.1.2</Version>
</PackageReference> </PackageReference>
</ItemGroup> </ItemGroup>
<ItemGroup> <ItemGroup>

View File

@ -61,7 +61,8 @@ namespace Wabbajack.Test
DownloadAndInstall(Game.SkyrimSpecialEdition, 12604, "SkyUI"), DownloadAndInstall(Game.SkyrimSpecialEdition, 12604, "SkyUI"),
DownloadAndInstall(Game.Fallout4, 11925, "Anti-Tank Rifle"), DownloadAndInstall(Game.Fallout4, 11925, "Anti-Tank Rifle"),
DownloadAndInstall(Game.SkyrimSpecialEdition, 4783, "Frost Armor UNP"), 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. // We're going to fully patch this mod from another source.
await modfiles[3].Download.DeleteAsync(); await modfiles[3].Download.DeleteAsync();

View File

@ -10,6 +10,7 @@ using OMODFramework;
using Wabbajack.Common.StatusFeed; using Wabbajack.Common.StatusFeed;
using Wabbajack.Common.StatusFeed.Errors; using Wabbajack.Common.StatusFeed.Errors;
using Wabbajack.Common; using Wabbajack.Common;
using Wabbajack.Common.FileSignatures;
using Utils = Wabbajack.Common.Utils; using Utils = Wabbajack.Common.Utils;
@ -17,19 +18,38 @@ namespace Wabbajack.VirtualFileSystem
{ {
public class FileExtractor 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<ExtractedFiles> ExtractAll(WorkQueue queue, AbsolutePath source, IEnumerable<RelativePath> OnlyFiles = null) public static async Task<ExtractedFiles> ExtractAll(WorkQueue queue, AbsolutePath source, IEnumerable<RelativePath> OnlyFiles = null)
{ {
try try
{ {
if (await BSADispatch.MightBeBSA(source)) var sig = await archiveSigs.MatchesAsync(source);
return await ExtractAllWithBSA(queue, source);
else if (source.Extension == Consts.OMOD) if (source.Extension == Consts.OMOD)
return await ExtractAllWithOMOD(source); return await ExtractAllWithOMOD(source);
else if (source.Extension == Consts.EXE)
return await ExtractAllExe(source); switch (sig)
else {
return await ExtractAllWith7Zip(source, OnlyFiles); 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) catch (Exception ex)
{ {