Add option to extract archives

ArcToolkitCLI/Program.cs

@@ -6,27 +6,36 @@ using System.Linq;
 using System.Text;
 using CommandLine;
 using COM3D2.Toolkit.Arc;
+using COM3D2.Toolkit.Arc.Files;
 using Ganss.IO;
 
 namespace ArcToolkitCLI
 {
-    internal interface IOptions
+    interface IOptions
     {
         [Value(0, MetaName = "input", HelpText = "Input ARC files")]
         IEnumerable<string> Input { get; set; }
     }
 
-    internal interface IDecryptionOptions
+    interface IOutputOptions
     {
-        [Option("arc-search-dir", Required = false, HelpText = "Directory where to search ARC file to decrypt the WARP file")]
+        [Option('o', "output", HelpText = "Output directory", Default = ".")]
+        string Output { get; set; }
+    }
+
+    interface IDecryptionOptions
+    {
+        [Option("arc-search-dir", Required = false,
+            HelpText = "Directory where to search ARC file to decrypt the WARP file")]
         string ArcDirectory { get; set; }
 
         [Option('k', "key", HelpText = "Decryption key as a base64 string (applied to all inputs)", Required = false)]
         string DecryptionKey { get; set; }
 
         [Option("key-file",
-                HelpText = "A file with decryption keys on each line. Format of the file is <decryption arc name>:<key in base64>.",
-                Required = false)]
+            HelpText =
+                "A file with decryption keys on each line. Format of the file is <decryption arc name>:<key in base64>.",
+            Required = false)]
         string KeyFile { get; set; }
 
         [Option("warc", Required = false, HelpText = "WARC file to use for decryption")]
@@ -34,44 +43,40 @@ namespace ArcToolkitCLI
     }
 
     [Verb("extract", HelpText = "Extract the contents of the given ARC files")]
-    internal class ExtractOptions : IOptions, IDecryptionOptions
+    class ExtractOptions : IOptions, IDecryptionOptions, IOutputOptions
     {
-        [Option('o', "output", HelpText = "Output directory", Required = true)]
-        public string Output { get; set; }
-
         public string ArcDirectory { get; set; }
         public string DecryptionKey { get; set; }
         public string KeyFile { get; set; }
 
         public string WarcFile { get; set; }
         public IEnumerable<string> Input { get; set; }
+        public string Output { get; set; }
     }
 
     [Verb("info", HelpText = "Display information about the given ARC files")]
-    internal class InfoOptions : IOptions
+    class InfoOptions : IOptions
     {
         [Option("only-key",
-                HelpText = "If the archive is a WARC file, output only the decryption key as a base64 string",
-                Required = false)]
+            HelpText = "If the archive is a WARC file, output only the decryption key as a base64 string",
+            Required = false)]
         public bool OnlyKey { get; set; }
 
         public IEnumerable<string> Input { get; set; }
     }
 
     [Verb("decrypt", HelpText = "Decrypts the provided WARP files")]
-    internal class DecryptOptions : IOptions, IDecryptionOptions
+    class DecryptOptions : IOptions, IDecryptionOptions, IOutputOptions
     {
-        [Option('o', "output", HelpText = "Output directory", Default = ".")]
-        public string Output { get; set; }
-
         public string ArcDirectory { get; set; }
         public string DecryptionKey { get; set; }
         public string KeyFile { get; set; }
         public string WarcFile { get; set; }
         public IEnumerable<string> Input { get; set; }
+        public string Output { get; set; }
     }
 
-    internal static class Errors
+    static class Errors
     {
         public enum ErrorCodes
         {
@@ -84,11 +89,11 @@ namespace ArcToolkitCLI
 
         public static readonly string[] ErrorMessages =
         {
-                "{0} is not a file",
-                "{0} is not a known ARC file",
-                "No key specified or the key file does not exist",
-                "The provided keyfile is not valid",
-                "The ARC {0} needs a key from {1}"
+            "{0} is not a file",
+            "{0} is not a known ARC file",
+            "No key specified or the key file does not exist",
+            "The provided keyfile is not valid",
+            "The ARC {0} needs a key from {1}"
         };
 
         public static int Error(ErrorCodes code, params object[] args)
@@ -98,13 +103,13 @@ namespace ArcToolkitCLI
         }
     }
 
-    internal static class Util
+    static class Util
     {
-        private static readonly char[] GlobCharacters = {'*', '{', '}', '[', ']', '?'};
+        static readonly char[] GlobCharacters = {'*', '{', '}', '[', ']', '?'};
 
         public static IEnumerable<FileSystemInfoBase> EnumerateFiles(IEnumerable<string> patterns)
         {
-            foreach (string pattern in patterns)
+            foreach (var pattern in patterns)
                 if (IsGlob(pattern))
                     foreach (var fileSystemInfoBase in Glob.Expand(pattern))
                         yield return fileSystemInfoBase;
@@ -112,26 +117,69 @@ namespace ArcToolkitCLI
                     yield return new FileInfoWrapper(new FileSystem(), new FileInfo(pattern));
         }
 
-        private static bool IsGlob(string path)
-        {
-            return path.IndexOfAny(GlobCharacters) >= 0;
-        }
+        static bool IsGlob(string path) { return path.IndexOfAny(GlobCharacters) >= 0; }
     }
 
-    internal class Program
+    class Program
     {
-        private static int Main(string[] args)
+        static int Decrypt(DecryptOptions opts)
         {
-            return Parser.Default.ParseArguments<InfoOptions, ExtractOptions, DecryptOptions>(args)
-                         .MapResult((InfoOptions opts) => DisplayInfo(opts),
-                                    (ExtractOptions opts) => Extract(opts),
-                                    (DecryptOptions opts) => Decrypt(opts),
-                                    errs => 1);
+            var errCode = GetDecryptionKeys(opts, out var keysDict);
+            if (errCode != 0)
+                return errCode;
+
+            Directory.CreateDirectory(opts.Output);
+
+            foreach (var file in Util.EnumerateFiles(opts.Input))
+            {
+                if (!file.Exists)
+                    return Errors.Error(Errors.ErrorCodes.NotAFile, file.Name);
+
+                using (var stream = File.OpenRead(file.FullName))
+                {
+                    var header = new byte[4];
+                    stream.Read(header, 0, header.Length);
+                    var headerString = Encoding.ASCII.GetString(header);
+                    stream.Position = 0;
+                    if (headerString == "warc")
+                        Console.WriteLine($"{file.Name} is a WARC file, skipping...");
+                    else if (headerString == "warp")
+                        try
+                        {
+                            var arcStream = WarpArc.DecryptWarp(stream, requestedFile =>
+                            {
+                                if (keysDict.TryGetValue("*", out var key))
+                                    return key;
+                                if (keysDict.TryGetValue(requestedFile, out key))
+                                    return key;
+                                if (Directory.Exists(opts.ArcDirectory) &&
+                                    File.Exists(Path.Combine(opts.ArcDirectory, requestedFile)))
+                                    return ReadKeyFromFile(Path.Combine(opts.ArcDirectory, requestedFile));
+                                if (File.Exists(opts.WarcFile))
+                                    return ReadKeyFromFile(opts.WarcFile);
+                                throw new FileNotFoundException("No key found for the requested ARC", requestedFile);
+                            });
+
+                            Console.WriteLine(
+                                $"Writing {Path.Combine(opts.Output, file.Name)} with length {arcStream.Length}");
+                            using (var output = File.Create(Path.Combine(opts.Output, file.Name)))
+                                arcStream.CopyTo(output);
+                        }
+                        catch (FileNotFoundException fe)
+                        {
+                            return Errors.Error(Errors.ErrorCodes.NoCorrectKey, file.Name, fe.FileName);
+                        }
+                    else
+                        return Errors.Error(Errors.ErrorCodes.UnknownArcType, file.Name);
+                }
+            }
+
+            return 0;
         }
 
-        private static int DisplayInfo(InfoOptions opts)
+        static int DisplayInfo(InfoOptions opts)
         {
-            bool first = true;
+            var first = true;
             foreach (var file in Util.EnumerateFiles(opts.Input))
             {
                 if (!file.Exists)
@@ -145,14 +193,14 @@ namespace ArcToolkitCLI
                 {
                     var header = new byte[4];
                     stream.Read(header, 0, header.Length);
-                    string headerString = Encoding.ASCII.GetString(header);
+                    var headerString = Encoding.ASCII.GetString(header);
                     stream.Position = 0;
 
                     if (headerString == "warc")
                     {
                         var key = new byte[2048];
                         stream.Read(key, 0, key.Length);
-                        string keyBase64 = Convert.ToBase64String(key);
+                        var keyBase64 = Convert.ToBase64String(key);
                         if (opts.OnlyKey)
                         {
                             Console.WriteLine($"{Path.GetFileName(file.FullName)}:{keyBase64}");
@@ -177,108 +225,135 @@ namespace ArcToolkitCLI
                         Console.WriteLine($"File name: {file.Name}");
                         Console.WriteLine("ARC type: WARP");
                         using (var br = new BinaryReader(stream))
-                            Console.WriteLine($"Needs a decryption key from the following ARC: {WarpArc.GetKeyWarpName(br)}");
+                            Console.WriteLine(
+                                $"Needs a decryption key from the following ARC: {WarpArc.GetKeyWarpName(br)}");
                     }
                     else
-                    {
                         return Errors.Error(Errors.ErrorCodes.UnknownArcType, file.Name);
-                    }
                 }
             }
 
             return 0;
         }
 
-        private static int Extract(ExtractOptions opts)
+        static int Extract(ExtractOptions opts)
         {
+            Dictionary<string, byte[]> keys = null;
+
+            Directory.CreateDirectory(opts.Output);
+
+            foreach (var file in Util.EnumerateFiles(opts.Input))
+            {
+                if (!file.Exists)
+                    return Errors.Error(Errors.ErrorCodes.NotAFile, file.Name);
+
+                using (var stream = File.OpenRead(file.FullName))
+                {
+                    var header = new byte[4];
+                    stream.Read(header, 0, header.Length);
+                    var headerString = Encoding.ASCII.GetString(header);
+                    stream.Position = 0;
+                    WarcArc arc = null;
+
+                    if (headerString == "warp")
+                    {
+                        if (keys == null)
+                        {
+                            var err = GetDecryptionKeys(opts, out keys);
+                            if (err != 0)
+                                return err;
+
+
+                            arc = new WarpArc(stream, requestedFile =>
+                            {
+                                if (keys.TryGetValue("*", out var key))
+                                    return key;
+                                if (keys.TryGetValue(requestedFile, out key))
+                                    return key;
+                                if (Directory.Exists(opts.ArcDirectory) &&
+                                    File.Exists(Path.Combine(opts.ArcDirectory, requestedFile)))
+                                    return ReadKeyFromFile(Path.Combine(opts.ArcDirectory, requestedFile));
+                                if (File.Exists(opts.WarcFile))
+                                    return ReadKeyFromFile(opts.WarcFile);
+                                throw new FileNotFoundException("No key found for the requested ARC", requestedFile);
+                            });
+                        }
+                    }
+                    else if (headerString == "warc")
+                        arc = new WarcArc(stream);
+
+                    if (arc == null)
+                        return Errors.Error(Errors.ErrorCodes.UnknownArcType, file.Name);
+
+                    Console.WriteLine($"Extracting {arc.Name}:");
+                    var arcDir = Path.Combine(opts.Output, Path.GetFileNameWithoutExtension(arc.Name));
+
+                    Directory.CreateDirectory(arcDir);
+                    foreach (var arcEntry in arc.Entries)
+                        ExtractArchive(arcEntry, arcDir);
+
+                    arc.Dispose();
+                }
+            }
+
             return 0;
         }
 
-        private static byte[] ReadKeyFromFile(string filename)
+        static void ExtractArchive(ArcEntry entry, string path)
         {
-            using (var br = new BinaryReader(File.OpenRead(filename)))
-                return br.ReadBytes(2048);
+            switch (entry)
+            {
+                case ArcFileEntry fe:
+                    var fullName = Path.Combine(path, fe.Name);
+                    Console.WriteLine(fullName);
+                    using (var file = File.Create(fullName))
+                        using (var data = fe.GetDataStream())
+                            data.CopyTo(file);
+                    break;
+                case ArcDirectoryEntry de:
+                    var subfolderPath = Path.Combine(path, de.Name);
+                    Directory.CreateDirectory(subfolderPath);
+                    foreach (var arcEntry in de.Children)
+                        ExtractArchive(arcEntry, subfolderPath);
+                    break;
+            }
         }
 
-        private static int Decrypt(DecryptOptions opts)
+        static int GetDecryptionKeys(IDecryptionOptions opts, out Dictionary<string, byte[]> keys)
         {
-            var keysDict = new Dictionary<string, byte[]>(StringComparer.InvariantCultureIgnoreCase);
+            keys = new Dictionary<string, byte[]>(StringComparer.InvariantCultureIgnoreCase);
             if (!string.IsNullOrWhiteSpace(opts.KeyFile))
             {
                 if (!File.Exists(opts.KeyFile))
                     return Errors.Error(Errors.ErrorCodes.KeyNotFound);
 
-                foreach (string line in File.ReadAllLines(opts.KeyFile))
+                foreach (var line in File.ReadAllLines(opts.KeyFile))
                 {
                     var parts = line.Split(':');
                     if (parts.Length != 2)
                         return Errors.Error(Errors.ErrorCodes.InvalidKeyFile);
-                    keysDict[parts[0].Trim()] = Convert.FromBase64String(parts[1].Trim());
+                    keys[parts[0].Trim()] = Convert.FromBase64String(parts[1].Trim());
                 }
             }
             else if (!string.IsNullOrWhiteSpace(opts.DecryptionKey))
-            {
-                keysDict["*"] = Convert.FromBase64String(opts.DecryptionKey);
-            }
+                keys["*"] = Convert.FromBase64String(opts.DecryptionKey);
             else if (string.IsNullOrWhiteSpace(opts.ArcDirectory) && string.IsNullOrWhiteSpace(opts.WarcFile))
-            {
                 return Errors.Error(Errors.ErrorCodes.KeyNotFound);
-            }
 
-            Directory.CreateDirectory(opts.Output);
-
-            foreach (var file in Util.EnumerateFiles(opts.Input))
-            {
-                if (!file.Exists)
-                    return Errors.Error(Errors.ErrorCodes.NotAFile, file.Name);
+            return 0;
+        }
 
-                using (var stream = File.OpenRead(file.FullName))
-                {
-                    var header = new byte[4];
-                    stream.Read(header, 0, header.Length);
-                    string headerString = Encoding.ASCII.GetString(header);
-                    stream.Position = 0;
-                    if (headerString == "warc")
-                    {
-                        Console.WriteLine($"{file.Name} is a WARC file, skipping...");
-                    }
-                    else if (headerString == "warp")
-                    {
-                        try
-                        {
-                            var arcStream = WarpArc.DecryptWarp(stream,
-                                                                requestedFile =>
-                                                                {
-                                                                    if (keysDict.TryGetValue("*", out var key))
-                                                                        return key;
-                                                                    if (keysDict.TryGetValue(requestedFile, out key))
-                                                                        return key;
-                                                                    if (Directory.Exists(opts.ArcDirectory)
-                                                                     && File.Exists(Path.Combine(opts.ArcDirectory, requestedFile)))
-                                                                        return ReadKeyFromFile(Path.Combine(opts.ArcDirectory, requestedFile));
-                                                                    if (File.Exists(opts.WarcFile))
-                                                                        return ReadKeyFromFile(opts.WarcFile);
-                                                                    throw new FileNotFoundException("No key found for the requested ARC",
-                                                                                                    requestedFile);
-                                                                });
-
-                            using (var output =
-                                    File.Create(Path.Combine(Path.GetDirectoryName(opts.Output), Path.GetFileNameWithoutExtension(file.Name))))
-                                arcStream.CopyTo(output);
-                        }
-                        catch (FileNotFoundException fe)
-                        {
-                            return Errors.Error(Errors.ErrorCodes.NoCorrectKey, file.Name, fe.FileName);
-                        }
-                    }
-                    else
-                    {
-                        return Errors.Error(Errors.ErrorCodes.UnknownArcType, file.Name);
-                    }
-                }
-            }
+        static int Main(string[] args)
+        {
+            return Parser.Default.ParseArguments<InfoOptions, ExtractOptions, DecryptOptions>(args).MapResult(
+                (InfoOptions opts) => DisplayInfo(opts), (ExtractOptions opts) => Extract(opts),
+                (DecryptOptions opts) => Decrypt(opts), errs => 1);
+        }
 
-            return 0;
+        static byte[] ReadKeyFromFile(string filename)
+        {
+            using (var br = new BinaryReader(File.OpenRead(filename)))
+                return br.ReadBytes(2048);
         }
     }
 }