From 61ff13f127132364021fec82c330edc3a2631c05 Mon Sep 17 00:00:00 2001
From: Razmoth <32140579+Razmoth@users.noreply.github.com>
Date: Tue, 6 Feb 2024 21:04:28 +0400
Subject: [PATCH] - [CLI] Added feature to search prebuild `AssetMap` to loaded
 filteded files only.

---
 AssetStudio.CLI/Program.cs  |  1 +
 AssetStudio/AssetsHelper.cs | 90 +++++++++++++++++++++++++++++++++++++
 AssetStudio/TypeFlags.cs    | 12 ++++-
 3 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/AssetStudio.CLI/Program.cs b/AssetStudio.CLI/Program.cs
index f5caf3c..7103be4 100644
--- a/AssetStudio.CLI/Program.cs
+++ b/AssetStudio.CLI/Program.cs
@@ -147,6 +147,7 @@ namespace AssetStudio.CLI
                 {
                     if (o.MapOp.HasFlag(MapOpType.Load))
                     {
+                        files = AssetsHelper.ParseAssetMap(o.MapName, o.MapType, classTypeFilter, o.NameFilter, o.ContainerFilter);
                     }
                     else
                     {
diff --git a/AssetStudio/AssetsHelper.cs b/AssetStudio/AssetsHelper.cs
index 844d48f..9e70e7f 100644
--- a/AssetStudio/AssetsHelper.cs
+++ b/AssetStudio/AssetsHelper.cs
@@ -505,6 +505,96 @@ namespace AssetStudio
             }
         }
 
+        public static string[] ParseAssetMap(string mapName, ExportListType mapType, ClassIDType[] typeFilter, Regex[] nameFilter, Regex[] containerFilter)
+        {
+            var matches = new HashSet<string>();
+
+            switch (mapType)
+            {
+                case ExportListType.MessagePack:
+                    {
+                        using var stream = File.OpenRead(mapName);
+                        var assetMap = MessagePackSerializer.Deserialize<AssetMap>(stream, MessagePackSerializerOptions.Standard.WithCompression(MessagePackCompression.Lz4BlockArray));
+                        foreach(var entry in assetMap.AssetEntries)
+                        {
+                            var isNameMatch = nameFilter.Length == 0 || nameFilter.Any(x => x.IsMatch(entry.Name));
+                            var isContainerMatch = containerFilter.Length == 0 || containerFilter.Any(x => x.IsMatch(entry.Container));
+                            var isTypeMatch = typeFilter.Length == 0 || typeFilter.Any(x => x == entry.Type);
+                            if (isNameMatch && isContainerMatch && isTypeMatch)
+                            {
+                                matches.Add(entry.Source);
+                            }
+                        }
+                    }
+
+                    break;
+                case ExportListType.XML:
+                    {
+                        using var stream = File.OpenRead(mapName);
+                        using var reader = XmlReader.Create(stream);
+                        reader.ReadToFollowing("Assets");
+                        reader.ReadToFollowing("Asset");
+                        do
+                        {
+                            reader.ReadToFollowing("Name");
+                            var name = reader.ReadInnerXml();
+
+                            var isNameMatch = nameFilter.Length == 0 || nameFilter.Any(x => x.IsMatch(name));
+
+                            reader.ReadToFollowing("Container");
+                            var container = reader.ReadInnerXml();
+
+                            var isContainerMatch = containerFilter.Length == 0 || containerFilter.Any(x => x.IsMatch(container));
+
+                            reader.ReadToFollowing("Type");
+                            var type = reader.ReadInnerXml();
+
+                            var isTypeMatch = typeFilter.Length == 0 || typeFilter.Any(x => x.ToString().Equals(type, StringComparison.OrdinalIgnoreCase));
+
+                            reader.ReadToFollowing("PathID");
+                            var pathID = reader.ReadInnerXml();
+
+                            reader.ReadToFollowing("Source");
+                            var source = reader.ReadInnerXml();
+
+                            if (isNameMatch && isContainerMatch && isTypeMatch)
+                            {
+                                matches.Add(source);
+                            }
+
+                            reader.ReadEndElement();
+                        } while (reader.ReadToNextSibling("Asset"));
+                    }
+
+                    break;
+                case ExportListType.JSON:
+                    {
+                        using var stream = File.OpenRead(mapName);
+                        using var file = new StreamReader(stream);
+                        using var reader = new JsonTextReader(file);
+
+                        var serializer = new JsonSerializer() { Formatting = Newtonsoft.Json.Formatting.Indented };
+                        serializer.Converters.Add(new StringEnumConverter());
+
+                        var entries = serializer.Deserialize<List<AssetEntry>>(reader);
+                        foreach (var entry in entries)
+                        {
+                            var isNameMatch = nameFilter.Length == 0 || nameFilter.Any(x => x.IsMatch(entry.Name));
+                            var isContainerMatch = containerFilter.Length == 0 || containerFilter.Any(x => x.IsMatch(entry.Container));
+                            var isTypeMatch = typeFilter.Length == 0 || typeFilter.Any(x => x == entry.Type);
+                            if (isNameMatch && isContainerMatch && isTypeMatch)
+                            {
+                                matches.Add(entry.Source);
+                            }
+                        }
+                    }
+
+                    break;
+            }
+
+            return matches.ToArray();
+        }
+
         private static void UpdateContainers(List<AssetEntry> assets, Game game)
         {
             if (game.Type.IsGISubGroup() && assets.Count > 0)
diff --git a/AssetStudio/TypeFlags.cs b/AssetStudio/TypeFlags.cs
index fda55ac..8d68598 100644
--- a/AssetStudio/TypeFlags.cs
+++ b/AssetStudio/TypeFlags.cs
@@ -1,4 +1,5 @@
-using System.Collections.Generic;
+using System;
+using System.Collections.Generic;
 
 namespace AssetStudio;
 public static class TypeFlags
@@ -44,3 +45,12 @@ public static class TypeFlags
         return false;
     }
 }
+
+[Flags]
+public enum TypeFlag
+{
+    None,
+    Parse,
+    Export,
+    Both = Parse | Export,
+}
\ No newline at end of file