From 1466be47c0268d2aabe17f5d3aa2dd43dd9f2699 Mon Sep 17 00:00:00 2001
From: Timothy Baldridge <tbaldridge@gmail.com>
Date: Sun, 21 Jul 2019 21:36:25 -0600
Subject: [PATCH] added patching support

---
 Wabbajack.Common/Data.cs  |  1 +
 Wabbajack.Common/Utils.cs |  9 +++++
 Wabbajack/Compiler.cs     | 69 ++++++++++++++++++++++++++++++++++++---
 3 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/Wabbajack.Common/Data.cs b/Wabbajack.Common/Data.cs
index 46f98712..c9653283 100644
--- a/Wabbajack.Common/Data.cs
+++ b/Wabbajack.Common/Data.cs
@@ -170,6 +170,7 @@ namespace Wabbajack.Common
         public dynamic IniData;
         public string Name;
         public string Meta;
+        public string AbsolutePath;
     }
 
     /// <summary>
diff --git a/Wabbajack.Common/Utils.cs b/Wabbajack.Common/Utils.cs
index 7b7ad49d..697bc6b3 100644
--- a/Wabbajack.Common/Utils.cs
+++ b/Wabbajack.Common/Utils.cs
@@ -120,5 +120,14 @@ namespace Wabbajack.Common
             }
         }
 
+        public static byte[] ReadAll(this Stream ins)
+        {
+            using (var ms = new MemoryStream())
+            {
+                ins.CopyTo(ms);
+                return ms.ToArray();
+            }
+        }
+
     }
 }
diff --git a/Wabbajack/Compiler.cs b/Wabbajack/Compiler.cs
index 6503545b..3c30c2d4 100644
--- a/Wabbajack/Compiler.cs
+++ b/Wabbajack/Compiler.cs
@@ -64,7 +64,7 @@ namespace Wabbajack
         public Action<string, long, long> Progress_Function { get; }
         public List<Directive> InstallDirectives { get; private set; }
         public List<Archive> SelectedArchives { get; private set; }
-
+        public List<RawSourceFile> AllFiles { get; private set; }
 
         public List<IndexedArchive> IndexedArchives;
 
@@ -111,6 +111,7 @@ namespace Wabbajack
             {
                 var info = metaname.FromJSON<IndexedArchive>();
                 info.Name = Path.GetFileName(file);
+                info.AbsolutePath = file;
 
                 var ini_name = file + ".meta";
                 if (ini_name.FileExists())
@@ -163,13 +164,13 @@ namespace Wabbajack
                                       .Where(p => p.FileExists())
                                       .Select(p => new RawSourceFile() { Path = Path.Combine(Consts.GameFolderFilesDir, p.RelativeTo(GamePath)), AbsolutePath = p });
 
-            var all_files = mo2_files.Concat(game_files).ToList();
+            AllFiles = mo2_files.Concat(game_files).ToList();
 
-            Info("Found {0} files to build into mod list", all_files.Count);
+            Info("Found {0} files to build into mod list", AllFiles.Count);
 
             var stack = MakeStack();
 
-            var results = all_files.AsParallel().Select(f => RunStack(stack, f)).ToList();
+            var results = AllFiles.AsParallel().Select(f => RunStack(stack, f)).ToList();
 
             var nomatch = results.OfType<NoMatch>();
             Info("No match for {0} files", nomatch.Count());
@@ -179,11 +180,71 @@ namespace Wabbajack
             InstallDirectives = results.Where(i => !(i is IgnoredDirectly)).ToList();
 
             GatherArchives();
+            BuildPatches();
 
             results.ToJSON("out.json");
 
         }
 
+
+        /// <summary>
+        /// Fills in the Patch fields in files that require them
+        /// </summary>
+        private void BuildPatches()
+        {
+            var groups = InstallDirectives.OfType<PatchedFromArchive>()
+                                          .GroupBy(p => p.ArchiveHash)
+                                          .ToList();
+
+            Info("Patching building patches from {0} archives", groups.Count);
+            var absolute_paths = AllFiles.ToDictionary(e => e.Path, e => e.AbsolutePath);
+            Parallel.ForEach(groups, group => BuildArchivePatches(group.Key, group, absolute_paths));
+
+            if (InstallDirectives.OfType<PatchedFromArchive>().FirstOrDefault(f => f.Patch == null) != null)
+            {
+                Error("Missing patches after generation, this should not happen");
+            }
+                  
+        }
+
+        private void BuildArchivePatches(string archive_sha, IEnumerable<PatchedFromArchive> group, Dictionary<string, string> absolute_paths)
+        {
+            var archive = IndexedArchives.First(a => a.Hash == archive_sha);
+            var paths = group.Select(g => g.From).ToHashSet();
+            var streams = new Dictionary<string, MemoryStream>();
+            Info("Etracting Patch Files from {0}", archive.Name);
+            // First we fetch the source files from the input archive
+            using (var a = new ArchiveFile(archive.AbsolutePath))
+            {
+                a.Extract(entry =>
+                {
+                if (!paths.Contains(entry.FileName)) return null;
+
+                var result = new MemoryStream();
+                streams.Add(entry.FileName, result);
+                return result;
+                }, false);
+            }
+
+            var extracted = streams.ToDictionary(k => k.Key, v => v.Value.ToArray());
+            // Now Create the patches
+            Parallel.ForEach(group, entry =>
+            {
+                Info("Patching {0}", entry.To);
+                var ss = extracted[entry.From];
+                using (var origin = new MemoryStream(ss))
+                using (var dest = File.OpenRead(absolute_paths[entry.To]))
+                using (var output = new MemoryStream())
+                {
+                    var a = origin.ReadAll();
+                    var b = dest.ReadAll();
+                    BSDiff.Create(a, b, output);
+                    entry.Patch = output.ToArray().ToBase64();
+                }
+            });
+
+        }
+
         private void GatherArchives()
         {
             var archives = IndexedArchives.GroupBy(a => a.Hash).ToDictionary(k => k.Key, k => k.First());