using System; using System.Collections.Concurrent; using System.Collections.Generic; using System.Diagnostics; using System.IO; using System.IO.Compression; using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.Logging; using Wabbajack.Common; using Wabbajack.Compiler.CompilationSteps; using Wabbajack.Downloaders; using Wabbajack.Downloaders.GameFile; using Wabbajack.DTOs; using Wabbajack.DTOs.Directives; using Wabbajack.DTOs.DownloadStates; using Wabbajack.DTOs.JsonConverters; using Wabbajack.FileExtractor.ExtractedFiles; using Wabbajack.Hashing.PHash; using Wabbajack.Hashing.xxHash64; using Wabbajack.Installer; using Wabbajack.Networking.WabbajackClientApi; using Wabbajack.Paths; using Wabbajack.Paths.IO; using Wabbajack.RateLimiter; using Wabbajack.VFS; namespace Wabbajack.Compiler; public abstract class ACompiler { protected readonly DownloadDispatcher _dispatcher; private readonly DTOSerializer _dtos; private readonly FileExtractor.FileExtractor _extractor; private readonly FileHashCache _hashCache; public readonly IGameLocator _locator; protected internal readonly ILogger _logger; private readonly TemporaryFileManager _manager; public readonly ParallelOptions _parallelOptions; public readonly IBinaryPatchCache _patchCache; private readonly AbsolutePath _stagingFolder; private readonly Stopwatch _updateStopWatch = new(); protected readonly Context _vfs; protected readonly Client _wjClient; public readonly IResource CompilerLimiter; private int _currentStep; private long _currentStepProgress; private long _maxStepProgress; public ConcurrentDictionary _patchOptions; public CompilerSettings _settings; public ConcurrentDictionary _sourceFileLinks; private string _statusText; private string _statusCategory; public List IndexedArchives = new(); public Dictionary> IndexedFiles = new(); public ModList ModList = new(); public AbsolutePath ModListImage; public ACompiler(ILogger logger, FileExtractor.FileExtractor extractor, FileHashCache hashCache, Context vfs, TemporaryFileManager manager, CompilerSettings settings, ParallelOptions parallelOptions, DownloadDispatcher dispatcher, Client wjClient, IGameLocator locator, DTOSerializer dtos, IResource compilerLimiter, IBinaryPatchCache patchCache, IImageLoader imageLoader) { CompilerLimiter = compilerLimiter; _logger = logger; _extractor = extractor; _hashCache = hashCache; _vfs = vfs; _manager = manager; _settings = settings; _stagingFolder = _manager.CreateFolder().Path; _parallelOptions = parallelOptions; _sourceFileLinks = new ConcurrentDictionary(); _dispatcher = dispatcher; _wjClient = wjClient; _locator = locator; _dtos = dtos; _patchOptions = new ConcurrentDictionary(); _sourceFileLinks = new ConcurrentDictionary(); _patchCache = patchCache; ImageLoader = imageLoader; _updateStopWatch = new Stopwatch(); } public IImageLoader ImageLoader { get; } protected long MaxSteps { get; set; } public CompilerSettings Settings { get => _settings; set => _settings = value; } public Dictionary> GameHashes { get; set; } = new(); public Dictionary GamesWithHashes { get; set; } = new(); public ILookup GameFiles { get; private set; } public bool IgnoreMissingFiles { get; set; } public List SelectedArchives { get; protected set; } = new(); public List InstallDirectives { get; protected set; } = new(); public List AllFiles { get; protected set; } = new(); public Dictionary ArchivesByFullPath { get; set; } = new(); public event EventHandler OnStatusUpdate; public void NextStep(string statusCategory, string statusText, long maxStepProgress = 1) { _updateStopWatch.Restart(); _maxStepProgress = maxStepProgress; _currentStep += 1; _statusText = statusText; _statusCategory = statusCategory; _logger.LogInformation("Compiler Step: {Step}", statusText); if (OnStatusUpdate != null) OnStatusUpdate(this, new StatusUpdate(statusCategory, statusText, Percent.FactoryPutInRange(_currentStep, MaxSteps), Percent.Zero, _currentStep)); } public void UpdateProgress(long stepProgress) { Interlocked.Add(ref _currentStepProgress, stepProgress); lock (_updateStopWatch) { if (_updateStopWatch.ElapsedMilliseconds < 100) return; _updateStopWatch.Restart(); } if (OnStatusUpdate != null) OnStatusUpdate(this, new StatusUpdate(_statusCategory, _statusText, Percent.FactoryPutInRange(_currentStep, MaxSteps), Percent.FactoryPutInRange(_currentStepProgress, _maxStepProgress), _currentStep)); } public void UpdateProgressAbsolute(long cur, long max) { _currentStepProgress = cur; _maxStepProgress = max; lock (_updateStopWatch) { if (_updateStopWatch.ElapsedMilliseconds < 100) return; _updateStopWatch.Restart(); } if (OnStatusUpdate != null) OnStatusUpdate(this, new StatusUpdate(_statusCategory, _statusText, Percent.FactoryPutInRange(_currentStep, MaxSteps), Percent.FactoryPutInRange(_currentStepProgress, _maxStepProgress), _currentStep)); } public abstract Task Begin(CancellationToken token); internal RelativePath IncludeId() { return Guid.NewGuid().ToString().ToRelativePath(); } internal async Task IncludeFile(byte[] data) { var id = IncludeId(); await _stagingFolder.Combine(id).WriteAllBytesAsync(data); return id; } internal async Task IncludeFile(Stream data) { var id = IncludeId(); await using var os = _stagingFolder.Combine(id).Open(FileMode.Create, FileAccess.Write); await data.CopyToAsync(os); return id; } internal AbsolutePath IncludeFile(out RelativePath id) { id = IncludeId(); return _stagingFolder.Combine(id); } internal async Task IncludeFile(string data) { var id = IncludeId(); await _stagingFolder.Combine(id).WriteAllTextAsync(data); return id; } internal async Task IncludeFile(Stream data, CancellationToken token) { var id = IncludeId(); await _stagingFolder.Combine(id).WriteAllAsync(data, token); return id; } internal async Task IncludeFile(AbsolutePath data, CancellationToken token) { await using var stream = data.Open(FileMode.Open); return await IncludeFile(stream, token); } internal async Task<(RelativePath, AbsolutePath)> IncludeString(string str) { var id = IncludeId(); var fullPath = _stagingFolder.Combine(id); await fullPath.WriteAllTextAsync(str); return (id, fullPath); } public async Task GatherMetaData() { _logger.LogInformation("Getting meta data for {count} archives", SelectedArchives.Count); NextStep("Building", "Gathering Metadata", SelectedArchives.Count); await SelectedArchives.PDoAll(CompilerLimiter, async a => { UpdateProgress(1); await _dispatcher.FillInMetadata(a); }); return true; } protected async Task IndexGameFileHashes() { NextStep("Compiling", "Indexing Game Files"); var gameFiles = new List(); if (_settings.UseGamePaths) { //taking the games in Settings.IncludedGames + currently compiling game so you can eg //include the stock game files if you are compiling for a VR game (ex: Skyrim + SkyrimVR) foreach (var ag in _settings.OtherGames.Append(_settings.Game).Distinct()) try { if (!_locator.TryFindLocation(ag, out var path)) { _logger.LogWarning("Game {game} was to be used in compilation but it is not installed", ag); return; } var mainFile = ag.MetaData().MainExecutable!.Value.RelativeTo(path); if (!mainFile.FileExists()) _logger.LogWarning("Main file {file} for {game} does not exist", mainFile, ag); var versionInfo = FileVersionInfo.GetVersionInfo(mainFile.ToString()); var files = await _wjClient.GetGameArchives(ag, versionInfo.FileVersion ?? "0.0.0.0"); gameFiles.AddRange(files); _logger.LogInformation($"Including {files.Length} stock game files from {ag} as download sources"); GameHashes[ag] = files.Select(f => f.Hash).ToHashSet(); IndexedArchives.AddRange(files.TryKeep(f => { var state = (GameFileSource) f.State; if (_vfs.Index.ByRootPath.TryGetValue(path.Combine(state.GameFile), out var vf)) { return (true, new IndexedArchive(vf) { Name = state.GameFile.ToString().Replace("/", "_").Replace("\\", "_"), State = state }); } return default; })); } catch (Exception e) { _logger.LogCritical(e, "Unable to find existing game files for {game}, skipping.", ag); } GamesWithHashes = GameHashes.SelectMany(g => g.Value.Select(h => (g, h))) .GroupBy(gh => gh.h) .ToDictionary(gh => gh.Key, gh => gh.Select(p => p.g.Key).ToArray()); } GameFiles = gameFiles.ToLookup(f => f.Hash); } protected async Task CleanInvalidArchivesAndFillState() { NextStep("Compiling", "Cleaning Invalid Archives", IndexedArchives.Count); var remove = await IndexedArchives.PMapAllBatchedAsync(CompilerLimiter, async a => { UpdateProgress(1); try { var resolved = await ResolveArchive(a); if (resolved == null) return a; a.State = resolved.State; return default; } catch (Exception ex) { _logger.LogWarning(ex, "While resolving archive {Archive}", a.Name); return a; } }).Where(x => x != null) .ToHashSet(); if (remove.Count == 0) return; _logger.LogWarning( "Removing {count} archives from the compilation state, this is probably not an issue but reference this if you have compilation failures", remove.Count); remove.Do(r => _logger.LogWarning("Resolution failed for: ({size} {hash}) {path}", r!.File.Size, r.File.Hash, r.File.FullPath)); IndexedArchives.RemoveAll(a => remove.Contains(a)); } protected async Task InferMetas(CancellationToken token) { async Task HasInvalidMeta(AbsolutePath filename) { var metaName = filename.WithExtension(Ext.Meta); if (!metaName.FileExists()) return true; try { var ini = metaName.LoadIniFile(); return await _dispatcher.ResolveArchive(ini["General"].ToDictionary(d => d.KeyName, d => d.Value)) == null; } catch (Exception e) { _logger.LogCritical(e, $"Exception while checking meta {filename}"); return true; } } var toFind = await _settings.Downloads.EnumerateFiles() .Where(f => f.Extension != Ext.Meta) .PMapAll(CompilerLimiter, async f => await HasInvalidMeta(f) ? f : default) .Where(f => f != default) .Where(f => f.FileExists()) .ToList(); NextStep("Initializing", "InferMetas", toFind.Count); if (toFind.Count == 0) return; _logger.LogInformation("Attempting to infer {count} metas from the server.", toFind.Count); await toFind.PDoAll(async f => { UpdateProgress(1); var vf = _vfs.Index.ByRootPath[f]; var archives = await _wjClient.GetArchivesForHash(vf.Hash); Archive? a = null; foreach (var archive in archives) if (await _dispatcher.Verify(archive, token)) { a = archive; break; } if (a == null) { await vf.AbsoluteName.WithExtension(Ext.Meta).WriteAllLinesAsync( new[] { "[General]", "unknownArchive=true" }, token); _logger.LogWarning("Could not infer meta for {archive} {hash}", f, vf.Hash); return; } _logger.LogInformation($"Inferred .meta for {vf.FullPath.FileName}, writing to disk"); await vf.AbsoluteName.WithExtension(Ext.Meta) .WriteAllTextAsync(_dispatcher.MetaIniSection(a), token); }); } protected async Task ExportModList(CancellationToken token) { NextStep("Finalizing", "Exporting Modlist"); _logger.LogInformation("Exporting ModList to {location}", _settings.OutputFile); // Modify readme and ModList image to relative paths if they exist if (_settings.ModListImage.FileExists()) ModList.Image = (RelativePath) "modlist-image.png"; await using (var of = _stagingFolder.Combine("modlist").Open(FileMode.Create, FileAccess.Write)) { await _dtos.Serialize(ModList, of); } await _wjClient.SendModListDefinition(ModList); _settings.OutputFile.Delete(); var allFiles = _stagingFolder.EnumerateFiles().ToList(); NextStep("Finalizing", "Writing Wabbajack File", allFiles.Count); await using (var fs = _settings.OutputFile.Open(FileMode.Create, FileAccess.Write)) { using var za = new ZipArchive(fs, ZipArchiveMode.Create); foreach (var f in allFiles) { UpdateProgress(1); var ze = za.CreateEntry((string) f.FileName); await using var os = ze.Open(); await using var ins = f.Open(FileMode.Open); await ins.CopyToAsync(os, token); } // Copy in modimage if (_settings.ModListImage.FileExists()) { var ze = za.CreateEntry((string) ModList.Image); await using var os = ze.Open(); await using var ins = _settings.ModListImage.Open(FileMode.Open); await ins.CopyToAsync(os, token); } } _logger.LogInformation("Exporting Modlist metadata"); var outputFileHash = await _hashCache.FileHashCachedAsync(_settings.OutputFile, token); if (outputFileHash == default) { _logger.LogCritical("Unable to hash Modlist Output File"); return; } var metadata = new DownloadMetadata { Size = _settings.OutputFile.Size(), Hash = outputFileHash, NumberOfArchives = ModList.Archives.Length, SizeOfArchives = ModList.Archives.Sum(a => a.Size), NumberOfInstalledFiles = ModList.Directives.Length, SizeOfInstalledFiles = ModList.Directives.Sum(a => a.Size) }; await using var metajson = _settings.OutputFile.WithExtension(new Extension(".meta.json")) .Open(FileMode.Create, FileAccess.Write); await _dtos.Serialize(metadata, metajson); _logger.LogInformation("Removing ModList staging folder"); _stagingFolder.DeleteDirectory(); } /// /// Fills in the Patch fields in files that require them /// protected async Task BuildPatches(CancellationToken token) { await using var tempPath = _manager.CreateFolder(); AbsolutePath TempPath(Hash file) { return tempPath.Path.Combine(file.ToHex()); } NextStep("Compiling","Looking for patches"); var toBuild = InstallDirectives.OfType() .Where(p => _patchOptions.GetValueOrDefault(p, Array.Empty()).Length > 0) .SelectMany(p => _patchOptions[p].Select(c => new PatchedFromArchive { To = p.To, Hash = p.Hash, FromHash = c.Hash, ArchiveHashPath = c.MakeRelativePaths(), Size = p.Size })) .ToArray(); if (toBuild.Any()) { NextStep("Compiling", "Generating Patches", toBuild.Length); var allFiles = (await toBuild.PMapAllBatchedAsync(CompilerLimiter, async f => { UpdateProgress(1); return new[] { _vfs.Index.FileForArchiveHashPath(f.ArchiveHashPath), FindDestFile(f.To) }; }).ToList()) .SelectMany(x => x) .DistinctBy(f => f.Hash) .ToHashSet(); _logger.LogInformation("Extracting {Count} ({Size}) files for building patches", allFiles.Count, allFiles.Sum(f => f.Size).ToFileSizeString()); NextStep("Compiling", "Extracting Patch Files", allFiles.Count); await _vfs.Extract(allFiles, async (vf, file) => { UpdateProgress(1); await using var ostream = TempPath(vf.Hash).Open(FileMode.Create, FileAccess.Write, FileShare.Read); await using var istream = await file.GetStream(); await istream.CopyToAsync(ostream, token); }, token); if (toBuild.Length == 0) return; NextStep("Compiling", "Generating Patch Files", toBuild.Length); await toBuild.PMapAllBatchedAsync(CompilerLimiter, async patch => { UpdateProgress(1); await using var src = TempPath(patch.FromHash).Open(FileMode.Open, FileAccess.Read, FileShare.Read); await using var dst = TempPath(patch.Hash).Open(FileMode.Open, FileAccess.Read, FileShare.Read); await _patchCache.CreatePatch(src, patch.FromHash, dst, patch.Hash); return patch; }).ToList(); } NextStep("Compiling", "Loading Patch Files"); // Load in the patches await InstallDirectives.OfType() .Where(p => p.PatchID == default) .PDoAll(CompilerLimiter, async pfa => { var patches = await _patchOptions[pfa] .SelectAsync(async c => (await _patchCache.GetPatch(c.Hash, pfa.Hash), c)) .ToList(); // Pick the best patch if (patches.All(p => p.Item1 != null)) { var (patch, file) = IncludePatches.PickPatch(this, patches); pfa.FromHash = file.Hash; pfa.ArchiveHashPath = file.MakeRelativePaths(); pfa.PatchID = await IncludeFile(await patch.cache.GetData(patch)); } }); var firstFailedPatch = InstallDirectives.OfType().FirstOrDefault(f => f.PatchID == default); if (firstFailedPatch != null) { _logger.LogCritical("Missing data from failed patch, starting data dump"); _logger.LogCritical("Dest File: {to}", firstFailedPatch.To); _logger.LogCritical("Options ({count}):", _patchOptions[firstFailedPatch].Length); foreach (var choice in _patchOptions[firstFailedPatch]) _logger.LogCritical(" {path}", choice.FullPath); _logger.LogCritical( "Missing patches after generation, this should not happen. First failure: {path}", firstFailedPatch.To); } } private VirtualFile FindDestFile(RelativePath to) { var abs = to.RelativeTo(_settings.Source); if (abs.FileExists()) return _vfs.Index.ByRootPath[abs]; if (to.InFolder(Consts.BSACreationDir)) { var bsaId = (RelativePath) ((string) to).Split('\\')[1]; var bsa = InstallDirectives.OfType().First(b => b.TempID == bsaId); var find = (RelativePath) Path.Combine(((string) to).Split('\\').Skip(2).ToArray()); return _vfs.Index.ByRootPath[_settings.Source.Combine(bsa.To)].Children.First(c => c.RelativeName == find); } throw new ArgumentException($"Couldn't load data for {to}"); } public async Task GenerateManifest() { NextStep("Finalizing", "Generating Manifest"); var manifest = new Manifest(ModList); await using var of = _settings.OutputFile.Open(FileMode.Create, FileAccess.Write); await _dtos.Serialize(manifest, of); } public async Task GatherArchives() { NextStep("Building", "Gathering Archives"); _logger.LogInformation("Building a list of archives based on the files required"); var hashes = InstallDirectives.OfType() .Select(a => a.ArchiveHashPath.Hash) .Distinct(); var archives = IndexedArchives.OrderByDescending(f => f.File.LastModified) .GroupBy(f => f.File.Hash) .ToDictionary(f => f.Key, f => f.First()); SelectedArchives.Clear(); SelectedArchives.AddRange(await hashes.PMapAll(CompilerLimiter, hash => { UpdateProgress(1); return ResolveArchive(hash, archives); }).ToList()); } public async Task ResolveArchive(Hash hash, IDictionary archives) { if (archives.TryGetValue(hash, out var found)) return (await ResolveArchive(found))!; throw new ArgumentException($"No match found for Archive sha: {hash.ToBase64()} this shouldn't happen"); } public async Task ResolveArchive(IndexedArchive archive) { if (archive.State == null && archive.IniData == null) { _logger.LogWarning( "No download metadata found for {archive}, please use MO2 to query info or add a .meta file and try again.", archive.Name); return null; } IDownloadState? state; if (archive.State == null) { state = await _dispatcher.ResolveArchive(archive.IniData!["General"] .ToDictionary(d => d.KeyName, d => d.Value)); if (state == null) { _logger.LogWarning("{archive} could not be handled by any of the downloaders", archive.Name); return null; } } else { state = archive.State; } var result = new Archive { State = state!, Name = archive.Name ?? "", Hash = archive.File.Hash, Size = archive.File.Size }; var token = new CancellationTokenSource(); token.CancelAfter(_settings.MaxVerificationTime); if (!await _dispatcher.Verify(result, token.Token)) _logger.LogWarning( "Unable to resolve link for {Archive}. If this is hosted on the Nexus the file may have been removed", result.State!.PrimaryKeyString); result.Meta = "[General]\n" + string.Join("\n", _dispatcher.MetaIni(result)); return result; } public async Task RunStack(IEnumerable stack, RawSourceFile source) { foreach (var step in stack) { var result = await step.Run(source); if (result != null) return result; } throw new InvalidDataException("Data fell out of the compilation stack"); } public abstract IEnumerable GetStack(); public abstract IEnumerable MakeStack(); public void PrintNoMatches(ICollection noMatches) { if (noMatches.Count > 0) foreach (var file in noMatches) _logger.LogWarning(" {fileTo} - {fileReason}", file.To, file.Reason); } protected async Task InlineFiles(CancellationToken token) { var grouped = ModList.Directives.OfType() .Where(f => f.SourceDataID == default) .GroupBy(f => _sourceFileLinks[f].File) .ToDictionary(k => k.Key); NextStep("Building", "Inlining Files"); if (grouped.Count == 0) return; await _vfs.Extract(grouped.Keys.ToHashSet(), async (vf, sfn) => { UpdateProgress(1); await using var stream = await sfn.GetStream(); var id = await IncludeFile(stream, token); foreach (var file in grouped[vf]) file.SourceDataID = id; }, token); } public bool CheckForNoMatchExit(ICollection noMatches) { if (noMatches.Count > 0) { _logger.LogCritical("Exiting due to no way to compile these files"); return true; } return false; } }