Implemented optional tables parameter

zbalkan · Feb 1, 2024 · e0bb5d0 · e0bb5d0
1 parent 354dbe5
commit e0bb5d0
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 5 deletions.
diff --git a/source/ditjson/Options.cs b/source/ditjson/Options.cs
@@ -1,4 +1,5 @@
-using CommandLine;
+using System.Collections.Generic;
+using CommandLine;
 
 namespace ditjson
 {
@@ -9,5 +10,8 @@ internal class Options
     {
         [Option('n', "ntds", Required = true, Default = "", HelpText = "Path to ntds.dit file")]
         public string Ntds { get; set; }
+
+        [Option('t', "tables", Required = false, Default = "", HelpText = "ntds.dit tables to include. Default: datatable, link_table")]
+        public IEnumerable<string> Tables { get; set; }
     }
 }
diff --git a/source/ditjson/Program.cs b/source/ditjson/Program.cs
@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
+using System.Linq;
 using System.Text;
 using System.Text.Json;
 using CommandLine;
@@ -20,6 +21,8 @@ internal static class Program
             WriteIndented = true,
         };
 
+        private static readonly string[] defaultTables = ["datatable", "link_table"];
+
         /// <summary>
         ///     Application entry point
         /// </summary>
@@ -61,11 +64,13 @@ internal static void RunOptions(Options opts)
             Api.JetAttachDatabase(session, opts.Ntds, AttachDatabaseGrbit.ReadOnly);
             Api.JetOpenDatabase(session, opts.Ntds, null, out var dbid, OpenDatabaseGrbit.ReadOnly);
 
-            var ntdsDictionary = new Dictionary<string, object>
+            var tablesToQuery = GetTablesToQuery(opts.Tables, session, dbid);
+
+            var ntdsDictionary = new Dictionary<string, object>();
+            foreach (var table in tablesToQuery)
             {
-                ["datatable"] = TableToList(session, dbid, "datatable"),
-                ["linktable"] = TableToList(session, dbid, "link_table")
-            };
+                ntdsDictionary.Add(table, TableToList(session, dbid, table));
+            }
 
             string json;
             try
@@ -87,6 +92,34 @@ internal static void RunOptions(Options opts)
             }
         }
 
+        private static List<string> GetTablesToQuery(IEnumerable<string>? tablesInOptions, Session session, JET_DBID dbid)
+        {
+            var tablesInDb = Api.GetTableNames(session, dbid);
+
+            List<string> tablesToQuery;
+
+            if (tablesInOptions == null)
+            {
+                // Add only the default tables
+                tablesToQuery = new List<string>(defaultTables);
+            }
+            else
+            {
+                // If user asks all
+                if (tablesInOptions.Count() == 1 && tablesInOptions.First().Equals("*", StringComparison.Ordinal))
+                {
+                    tablesToQuery = new List<string>(tablesInDb);
+                }
+                else
+                {
+                    // if user asks oly specific tables
+                    tablesToQuery = new List<string>(tablesInOptions.Where(t => tablesInDb.Contains(t)));
+                }
+            }
+
+            return tablesToQuery;
+        }
+
         /// <summary>
         /// Return the string format of a byte array.
         /// </summary>