using System; using System.Collections.Generic; using System.Collections.Immutable; using System.IO; using System.Linq; using System.Reactive.Linq; using System.Reactive.Subjects; using System.Text; using System.Threading.Tasks; using Alphaleonis.Win32.Filesystem; using Wabbajack.Common; using Wabbajack.Common.CSP; using Directory = Alphaleonis.Win32.Filesystem.Directory; using File = System.IO.File; using FileInfo = Alphaleonis.Win32.Filesystem.FileInfo; using Path = Alphaleonis.Win32.Filesystem.Path; namespace Wabbajack.VirtualFileSystem { public class Context { static Context() { Utils.Log("Cleaning VFS, this may take a bit of time"); Utils.DeleteDirectory(StagingFolder); } public const ulong FileVersion = 0x03; public const string Magic = "WABBAJACK VFS FILE"; private static readonly AbsolutePath StagingFolder = ((RelativePath)"vfs_staging").RelativeToWorkingDirectory(); public IndexRoot Index { get; private set; } = IndexRoot.Empty; /// /// A stream of tuples of ("Update Title", 0.25) which represent the name of the current task /// and the current progress. /// public IObservable<(string, float)> ProgressUpdates => _progressUpdates; private readonly Subject<(string, float)> _progressUpdates = new Subject<(string, float)>(); public StatusUpdateTracker UpdateTracker { get; set; } = new StatusUpdateTracker(1); public WorkQueue Queue { get; } public bool UseExtendedHashes { get; set; } public Context(WorkQueue queue, bool extendedHashes = false) { Queue = queue; UseExtendedHashes = extendedHashes; } public static TemporaryDirectory GetTemporaryFolder() { return new TemporaryDirectory(((RelativePath)Guid.NewGuid().ToString()).RelativeTo(StagingFolder)); } public async Task AddRoot(AbsolutePath root) { var filtered = Index.AllFiles.Where(file => file.IsNative && ((AbsolutePath) file.Name).Exists).ToList(); var byPath = filtered.ToImmutableDictionary(f => f.Name); var filesToIndex = root.EnumerateFiles().Distinct().ToList(); var results = Channel.Create(1024, ProgressUpdater($"Indexing {root}", filesToIndex.Count)); var allFiles = await filesToIndex .PMap(Queue, async f => { if (byPath.TryGetValue(f, out var found)) { if (found.LastModified == f.LastModifiedUtc.AsUnixTime() && found.Size == f.Size) return found; } return await VirtualFile.Analyze(this, null, f, f, true); }); var newIndex = await IndexRoot.Empty.Integrate(filtered.Concat(allFiles).ToList()); lock (this) { Index = newIndex; } return newIndex; } public async Task AddRoots(List roots) { var native = Index.AllFiles.Where(file => file.IsNative).ToDictionary(file => file.StagedPath); var filtered = Index.AllFiles.Where(file => ((AbsolutePath)file.Name).Exists).ToList(); var filesToIndex = roots.SelectMany(root => root.EnumerateFiles()).ToList(); var allFiles = await filesToIndex .PMap(Queue, async f => { Utils.Status($"Indexing {Path.GetFileName((string)f)}"); if (native.TryGetValue(f, out var found)) { if (found.LastModified == f.LastModifiedUtc.AsUnixTime() && found.Size == f.Size) return found; } return await VirtualFile.Analyze(this, null, f, f, true); }); var newIndex = await IndexRoot.Empty.Integrate(filtered.Concat(allFiles).ToList()); lock (this) { Index = newIndex; } return newIndex; } class Box { public T Value { get; set; } } private Func, IObservable> ProgressUpdater(string s, float totalCount) { if (totalCount == 0) totalCount = 1; var box = new Box(); return sub => sub.Select(itm => { box.Value += 1; _progressUpdates.OnNext((s, box.Value / totalCount)); return itm; }); } public async Task WriteToFile(AbsolutePath filename) { await using var fs = filename.Create(); await using var bw = new BinaryWriter(fs, Encoding.UTF8, true); fs.SetLength(0); bw.Write(Encoding.ASCII.GetBytes(Magic)); bw.Write(FileVersion); bw.Write((ulong) Index.AllFiles.Count); (await Index.AllFiles .PMap(Queue, f => { var ms = new MemoryStream(); f.Write(ms); return ms; })) .Do(ms => { var size = ms.Position; ms.Position = 0; bw.Write((ulong) size); ms.CopyTo(fs); }); Utils.Log($"Wrote {fs.Position.ToFileSizeString()} file as vfs cache file {filename}"); } public async Task IntegrateFromFile(AbsolutePath filename) { try { await using var fs = filename.OpenRead(); using var br = new BinaryReader(fs, Encoding.UTF8, true); var magic = Encoding.ASCII.GetString(br.ReadBytes(Encoding.ASCII.GetBytes(Magic).Length)); var fileVersion = br.ReadUInt64(); if (fileVersion != FileVersion || magic != Magic) throw new InvalidDataException("Bad Data Format"); var numFiles = br.ReadUInt64(); var files = Enumerable.Range(0, (int) numFiles) .Select(idx => { var size = br.ReadUInt64(); var bytes = new byte[size]; br.BaseStream.Read(bytes, 0, (int) size); return VirtualFile.Read(this, bytes); }).ToList(); var newIndex = await Index.Integrate(files); lock (this) { Index = newIndex; } } catch (IOException) { filename.Delete(); } } public async Task Stage(IEnumerable files) { var grouped = files.SelectMany(f => f.FilesInFullPath) .Distinct() .Where(f => f.Parent != null) .GroupBy(f => f.Parent) .OrderBy(f => f.Key?.NestingFactor ?? 0) .ToList(); var paths = new List(); foreach (var group in grouped) { var tmpPath = ((RelativePath)Guid.NewGuid().ToString()).RelativeTo(StagingFolder); await FileExtractor.ExtractAll(Queue, group.Key.StagedPath, tmpPath); paths.Add(tmpPath); foreach (var file in group) file.StagedPath = file.RelativeName.RelativeTo(tmpPath); } return () => { paths.Do(p => { p.DeleteDirectory(); }); }; } public List GetPortableState(IEnumerable files) { return files.SelectMany(f => f.FilesInFullPath) .Distinct() .Select(f => new PortableFile { Name = f.Parent != null ? f.Name : null, Hash = f.Hash, ParentHash = f.Parent?.Hash ?? Hash.Empty, Size = f.Size }).ToList(); } public async Task IntegrateFromPortable(List state, Dictionary links) { var indexedState = state.GroupBy(f => f.ParentHash) .ToDictionary(f => f.Key, f => (IEnumerable) f); var parents = await indexedState[Hash.Empty] .PMap(Queue,f => VirtualFile.CreateFromPortable(this, indexedState, links, f)); var newIndex = await Index.Integrate(parents); lock (this) { Index = newIndex; } } public async Task> StageWith(IEnumerable files) { return new DisposableList(await Stage(files), files); } #region KnownFiles private List _knownFiles = new List(); public void AddKnown(IEnumerable known) { _knownFiles.AddRange(known); } public async Task BackfillMissing() { var newFiles = _knownFiles.Where(f => f.Paths.Length == 0) .GroupBy(f => f.BaseHash) .ToDictionary(f => f.Key, s => new VirtualFile() { Name = s.First().Paths[0], Hash = s.First().BaseHash, Context = this }); var parentchild = new Dictionary<(VirtualFile, RelativePath), VirtualFile>(); void BackFillOne(HashRelativePath file) { var parent = newFiles[file.BaseHash]; foreach (var path in file.Paths.Skip(1)) { if (parentchild.TryGetValue((parent, path), out var foundParent)) { parent = foundParent; continue; } var nf = new VirtualFile {Name = path, Parent = parent}; parent.Children = parent.Children.Add(nf); parentchild.Add((parent, path), nf); parent = nf; } } _knownFiles.Where(f => f.Paths.Length > 1).Do(BackFillOne); var newIndex = await Index.Integrate(newFiles.Values.ToList()); lock (this) Index = newIndex; _knownFiles = new List(); } #endregion } public class DisposableList : List, IDisposable { private Action _unstage; public DisposableList(Action unstage, IEnumerable files) : base(files) { _unstage = unstage; } public void Dispose() { _unstage(); } } public class IndexRoot { public static IndexRoot Empty = new IndexRoot(); public IndexRoot(ImmutableList aFiles, ImmutableDictionary byFullPath, ImmutableDictionary> byHash, ImmutableDictionary byRoot, ImmutableDictionary> byName) { AllFiles = aFiles; ByFullPath = byFullPath; ByHash = byHash; ByRootPath = byRoot; ByName = byName; } public IndexRoot() { AllFiles = ImmutableList.Empty; ByFullPath = ImmutableDictionary.Empty; ByHash = ImmutableDictionary>.Empty; ByRootPath = ImmutableDictionary.Empty; ByName = ImmutableDictionary>.Empty; } public ImmutableList AllFiles { get; } public ImmutableDictionary ByFullPath { get; } public ImmutableDictionary> ByHash { get; } public ImmutableDictionary> ByName { get; set; } public ImmutableDictionary ByRootPath { get; } public async Task Integrate(ICollection files) { Utils.Log($"Integrating {files.Count} files"); var allFiles = AllFiles.Concat(files).GroupBy(f => f.Name).Select(g => g.Last()).ToImmutableList(); var byFullPath = Task.Run(() => allFiles.SelectMany(f => f.ThisAndAllChildren) .ToImmutableDictionary(f => f.FullPath)); var byHash = Task.Run(() => allFiles.SelectMany(f => f.ThisAndAllChildren) .Where(f => f.Hash != Hash.Empty) .ToGroupedImmutableDictionary(f => f.Hash)); var byName = Task.Run(() => allFiles.SelectMany(f => f.ThisAndAllChildren) .ToGroupedImmutableDictionary(f => f.Name)); var byRootPath = Task.Run(() => allFiles.ToImmutableDictionary(f => f.Name as AbsolutePath)); var result = new IndexRoot(allFiles, await byFullPath, await byHash, await byRootPath, await byName); Utils.Log($"Done integrating"); return result; } public VirtualFile FileForArchiveHashPath(HashRelativePath argArchiveHashPath) { var cur = ByHash[argArchiveHashPath.BaseHash].First(f => f.Parent == null); return argArchiveHashPath.Paths.Aggregate(cur, (current, itm) => ByName[itm].First(f => f.Parent == current)); } } public class TemporaryDirectory : IDisposable { public TemporaryDirectory(AbsolutePath name) { FullName = name; if (!FullName.Exists) FullName.CreateDirectory(); } public AbsolutePath FullName { get; } public void Dispose() { if (FullName.Exists) Utils.DeleteDirectory(FullName); } } }