wabbajack/Wabbajack.VFS/VirtualFile.cs

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using Wabbajack.Common;
using Wabbajack.Common.FileSignatures;
using Wabbajack.DTOs.Streams;
using Wabbajack.DTOs.Texture;
using Wabbajack.Hashing.PHash;
using Wabbajack.Hashing.xxHash64;
using Wabbajack.Paths;
using Wabbajack.Paths.IO;

namespace Wabbajack.VFS
{
    public class VirtualFile
    {
        private static readonly HashSet<Extension> TextureExtensions = new()
            { new Extension(".dds"), new Extension(".tga") };

        private static readonly SignatureChecker DDSSig = new(FileType.DDS);

        private IEnumerable<VirtualFile> _thisAndAllChildren;

        public IPath Name { get; internal set; }

        public RelativePath RelativeName => (RelativePath)Name;

        public AbsolutePath AbsoluteName => (AbsolutePath)Name;


        public FullPath FullPath { get; private set; }

        public Hash Hash { get; internal set; }
        public ImageState? ImageState { get; internal set; }
        public long Size { get; internal set; }

        public ulong LastModified { get; internal set; }

        public ulong LastAnalyzed { get; internal set; }

        public VirtualFile Parent { get; internal set; }

        public Context Context { get; set; }

        /// <summary>
        ///     Returns the nesting factor for this file. Native files will have a nesting of 1, the factor
        ///     goes up for each nesting of a file in an archive.
        /// </summary>
        public int NestingFactor
        {
            get
            {
                var cnt = 0;
                var cur = this;
                while (cur != null)
                {
                    cnt += 1;
                    cur = cur.Parent;
                }

                return cnt;
            }
        }

        public ImmutableList<VirtualFile> Children { get; internal set; } = ImmutableList<VirtualFile>.Empty;

        public bool IsArchive => Children != null && Children.Count > 0;

        public bool IsNative => Parent == null;

        public IEnumerable<VirtualFile> ThisAndAllChildren
        {
            get
            {
                if (_thisAndAllChildren == null)
                    _thisAndAllChildren = Children.SelectMany(child => child.ThisAndAllChildren).Append(this).ToList();

                return _thisAndAllChildren;
            }
        }


        /// <summary>
        ///     Returns all the virtual files in the path to this file, starting from the root file.
        /// </summary>
        public IEnumerable<VirtualFile> FilesInFullPath
        {
            get
            {
                var stack = ImmutableStack<VirtualFile>.Empty;
                var cur = this;
                while (cur != null)
                {
                    stack = stack.Push(cur);
                    cur = cur.Parent;
                }

                return stack;
            }
        }


        public VirtualFile TopParent => IsNative ? this : Parent.TopParent;


        public T ThisAndAllChildrenReduced<T>(T acc, Func<T, VirtualFile, T> fn)
        {
            acc = fn(acc, this);
            return Children.Aggregate(acc, (current, itm) => itm.ThisAndAllChildrenReduced(current, fn));
        }

        public void ThisAndAllChildrenReduced(Action<VirtualFile> fn)
        {
            fn(this);
            foreach (var itm in Children)
                itm.ThisAndAllChildrenReduced(fn);
        }

        private static VirtualFile ConvertFromIndexedFile(Context context, IndexedVirtualFile file, IPath path,
            VirtualFile vparent, IStreamFactory extractedFile)
        {
            var vself = new VirtualFile
            {
                Context = context,
                Name = path,
                Parent = vparent,
                Size = file.Size,
                LastModified = extractedFile.LastModifiedUtc.AsUnixTime(),
                LastAnalyzed = DateTime.Now.AsUnixTime(),
                Hash = file.Hash,
                ImageState = file.ImageState
            };

            vself.FillFullPath();

            vself.Children = file.Children.Select(f => ConvertFromIndexedFile(context, f, f.Name, vself, extractedFile))
                .ToImmutableList();

            return vself;
        }


        internal IndexedVirtualFile ToIndexedVirtualFile()
        {
            return new IndexedVirtualFile
            {
                Hash = Hash,
                ImageState = ImageState,
                Name = Name,
                Children = Children.Select(c => c.ToIndexedVirtualFile()).ToList(),
                Size = Size
            };
        }

        public static async Task<VirtualFile> Analyze(Context context, VirtualFile? parent,
            IStreamFactory extractedFile,
            IPath relPath, CancellationToken token, int depth = 0)
        {
            Hash hash;
            using (var job = await context.Limiter.Begin("Hash file", 0, token))
            {
                if (extractedFile is NativeFileStreamFactory)
                {
                    var absPath = (AbsolutePath)extractedFile.Name;
                    job.Size = absPath.Size();
                    hash = await context.HashCache.FileHashCachedAsync(absPath, token, job);
                }
                else
                {
                    await using var hstream = await extractedFile.GetStream();
                    job.Size = hstream.Length;
                    hash = await hstream.HashingCopy(Stream.Null, token, job);
                }
            }

            if (context.VfsCache.TryGetFromCache(context, parent, relPath, extractedFile, hash, out var vself))
                return vself;


            await using var stream = await extractedFile.GetStream();
            var sig = await FileExtractor.FileExtractor.ArchiveSigs.MatchesAsync(stream);
            stream.Position = 0;

            var self = new VirtualFile
            {
                Context = context,
                Name = relPath,
                Parent = parent,
                Size = stream.Length,
                LastModified = extractedFile.LastModifiedUtc.AsUnixTime(),
                LastAnalyzed = DateTime.Now.AsUnixTime(),
                Hash = hash
            };


            if (TextureExtensions.Contains(relPath.FileName.Extension) && await DDSSig.MatchesAsync(stream) != null)
                try
                {
                    using var job = await context.Limiter.Begin("Perceptual hash", self.Size, token);
                    self.ImageState = await ImageLoader.Load(stream);
                    await job.Report((int)self.Size, token);
                    stream.Position = 0;
                }
                catch (Exception)
                {
                }

            self.FillFullPath(depth);


            // Can't extract, so return
            if (!sig.HasValue ||
                !FileExtractor.FileExtractor.ExtractableExtensions.Contains(relPath.FileName.Extension))
            {
                await context.VfsCache.WriteToCache(self);
                return self;
            }

            try
            {
                var list = await context.Extractor.GatheringExtract(extractedFile,
                    _ => true,
                    async (path, sfactory) => await Analyze(context, self, sfactory, path, token, depth + 1),
                    token);

                self.Children = list.Values.ToImmutableList();
            }
            catch (EndOfStreamException)
            {
                return self;
            }
            catch (Exception ex)
            {
                context.Logger.LogError(ex, "Error while examining the contents of {path}", relPath.FileName);
                throw;
            }

            await context.VfsCache.WriteToCache(self);
            return self;
        }


        internal void FillFullPath()
        {
            var depth = 0;
            var self = this;
            while (self.Parent != null)
            {
                depth += 1;
                self = self.Parent;
            }

            FillFullPath(depth);
        }

        internal void FillFullPath(int depth)
        {
            if (depth == 0)
            {
                FullPath = new FullPath((AbsolutePath)Name);
            }
            else
            {
                var paths = new RelativePath[depth];
                var self = this;
                for (var idx = depth; idx != 0; idx -= 1)
                {
                    paths[idx - 1] = self.RelativeName;
                    self = self.Parent;
                }

                FullPath = new FullPath(self.AbsoluteName, paths);
            }
        }

        public void Write(BinaryWriter bw)
        {
            bw.Write(Name.ToString() ?? string.Empty);
            bw.Write(Size);
            bw.Write(LastModified);
            bw.Write(LastModified);
            bw.Write((ulong)Hash);
            bw.Write(Children.Count);
            foreach (var child in Children)
                child.Write(bw);
        }

        public static VirtualFile Read(Context context, byte[] data)
        {
            using var ms = new MemoryStream(data);
            using var br = new BinaryReader(ms);
            return Read(context, null, br);
        }

        private static VirtualFile Read(Context context, VirtualFile parent, BinaryReader br)
        {
            var vf = new VirtualFile
            {
                Name = br.ReadIPath(),
                Size = br.ReadInt64(),
                LastModified = br.ReadUInt64(),
                LastAnalyzed = br.ReadUInt64(),
                Hash = Hash.FromULong(br.ReadUInt64()),
                Context = context,
                Parent = parent,
                Children = ImmutableList<VirtualFile>.Empty
            };
            vf.FullPath = new FullPath(vf.AbsoluteName);
            var children = br.ReadInt32();
            for (var i = 0; i < children; i++)
            {
                var child = Read(context, vf, br, (AbsolutePath)vf.Name, new RelativePath[0]);
                vf.Children = vf.Children.Add(child);
            }

            return vf;
        }

        private static VirtualFile Read(Context context, VirtualFile parent, BinaryReader br, AbsolutePath top,
            RelativePath[] subpaths)
        {
            var name = (RelativePath)br.ReadIPath();
            subpaths = subpaths.Add(name);
            var vf = new VirtualFile
            {
                Name = name,
                Size = br.ReadInt64(),
                LastModified = br.ReadUInt64(),
                LastAnalyzed = br.ReadUInt64(),
                Hash = Hash.FromULong(br.ReadUInt64()),
                Context = context,
                Parent = parent,
                Children = ImmutableList<VirtualFile>.Empty,
                FullPath = new FullPath(top, subpaths)
            };

            var children = br.ReadInt32();
            for (var i = 0; i < children; i++)
            {
                var child = Read(context, vf, br, top, subpaths);
                vf.Children = vf.Children.Add(child);
            }

            return vf;
        }

        public HashRelativePath MakeRelativePaths()
        {
            var paths = new RelativePath[FilesInFullPath.Count() - 1];

            var idx = 0;
            foreach (var itm in FilesInFullPath.Skip(1))
            {
                paths[idx] = (RelativePath)itm.Name;
                idx += 1;
            }

            var path = new HashRelativePath(FilesInFullPath.First().Hash, paths);
            return path;
        }

        public VirtualFile InSameFolder(RelativePath relativePath)
        {
            var newPath = FullPath.InSameFolder(relativePath);
            return Context.Index.ByFullPath.TryGetValue(newPath, out var found) ? found : null;
        }
    }
}