From c73493b505fa005fc9dff6417e2b5bfcd419e96d Mon Sep 17 00:00:00 2001
From: Alex Barney <thealexbarney@gmail.com>
Date: Fri, 15 Mar 2019 21:12:43 -0500
Subject: [PATCH] hactoolnet: Use new NCA APIs

---
 src/LibHac/IO/NcaUtils/Nca.cs           | 15 ++++++-----
 src/LibHac/IO/NcaUtils/NcaExtensions.cs | 17 ++++++++++---
 src/LibHac/IO/NcaUtils/NcaStructs.cs    |  7 ------
 src/hactoolnet/ProcessNca.cs            | 33 ++++++++-----------------
 src/hactoolnet/ProcessRomfs.cs          |  8 +++---
 src/hactoolnet/ProcessSwitchFs.cs       | 23 ++++-------------
 src/hactoolnet/ProcessXci.cs            | 23 ++++-------------
 7 files changed, 47 insertions(+), 79 deletions(-)

diff --git a/src/LibHac/IO/NcaUtils/Nca.cs b/src/LibHac/IO/NcaUtils/Nca.cs
index 5f4c7118..be01ded8 100644
--- a/src/LibHac/IO/NcaUtils/Nca.cs
+++ b/src/LibHac/IO/NcaUtils/Nca.cs
@@ -10,7 +10,7 @@ namespace LibHac.IO.NcaUtils
     {
         private const int HeaderSize = 0xc00;
         private const int HeaderSectorSize = 0x200;
-        
+
         public NcaHeader Header { get; }
         public string NcaId { get; set; }
         public string Filename { get; set; }
@@ -91,6 +91,11 @@ namespace LibHac.IO.NcaUtils
             return sect.Header.EncryptionType == NcaEncryptionType.None || !IsMissingTitleKey && string.IsNullOrWhiteSpace(MissingKeyName);
         }
 
+        public bool CanOpenSection(NcaSectionType type)
+        {
+            return CanOpenSection(GetSectionIndexFromType(type));
+        }
+
         private IStorage OpenEncryptedStorage(int index)
         {
             if (index < 0 || index > 3) throw new ArgumentOutOfRangeException(nameof(index));
@@ -275,11 +280,9 @@ namespace LibHac.IO.NcaUtils
             return Sections[index] != null;
         }
 
-        public bool SectionIsDecryptable(int index)
+        public bool SectionExists(NcaSectionType type)
         {
-            if (!SectionExists(index)) return false;
-
-            return Sections[index].Header.EncryptionType == NcaEncryptionType.None || !IsMissingTitleKey && string.IsNullOrWhiteSpace(MissingKeyName);
+            return SectionExists(GetSectionIndexFromType(type));
         }
 
         /// <summary>
@@ -363,7 +366,7 @@ namespace LibHac.IO.NcaUtils
         {
             if (Header != null)
             {
-               return Header.Version;
+                return Header.Version;
             }
             else
             {
diff --git a/src/LibHac/IO/NcaUtils/NcaExtensions.cs b/src/LibHac/IO/NcaUtils/NcaExtensions.cs
index 115a8c0b..c7c04bc2 100644
--- a/src/LibHac/IO/NcaUtils/NcaExtensions.cs
+++ b/src/LibHac/IO/NcaUtils/NcaExtensions.cs
@@ -25,16 +25,27 @@ namespace LibHac.IO.NcaUtils
                 .WriteAllBytes(filename, logger);
         }
 
+        public static void ExportSection(this Nca nca, NcaSectionType type, string filename, bool raw = false,
+            IntegrityCheckLevel integrityCheckLevel = IntegrityCheckLevel.None, IProgressReport logger = null)
+        {
+            nca.OpenStorage(type, integrityCheckLevel, raw)
+                .WriteAllBytes(filename, logger);
+        }
+
         public static void ExtractSection(this Nca nca, int index, string outputDir,
             IntegrityCheckLevel integrityCheckLevel = IntegrityCheckLevel.None, IProgressReport logger = null)
         {
-            if (index < 0 || index > 3) throw new IndexOutOfRangeException();
-            if (!nca.SectionIsDecryptable(index)) return;
-
             IFileSystem fs = nca.OpenFileSystem(index, integrityCheckLevel);
             fs.Extract(outputDir, logger);
         }
 
+        public static void ExtractSection(this Nca nca, NcaSectionType type, string outputDir,
+            IntegrityCheckLevel integrityCheckLevel = IntegrityCheckLevel.None, IProgressReport logger = null)
+        {
+            IFileSystem fs = nca.OpenFileSystem(type, integrityCheckLevel);
+            fs.Extract(outputDir, logger);
+        }
+
         public static Validity VerifyNca(this Nca nca, IProgressReport logger = null, bool quiet = false)
         {
             for (int i = 0; i < 3; i++)
diff --git a/src/LibHac/IO/NcaUtils/NcaStructs.cs b/src/LibHac/IO/NcaUtils/NcaStructs.cs
index 07d6a731..df570ee7 100644
--- a/src/LibHac/IO/NcaUtils/NcaStructs.cs
+++ b/src/LibHac/IO/NcaUtils/NcaStructs.cs
@@ -193,13 +193,6 @@ namespace LibHac.IO.NcaUtils
         }
     }
 
-    public enum ProgramPartitionType
-    {
-        Code,
-        Data,
-        Logo
-    };
-
     public enum NcaSectionType
     {
         Code,
diff --git a/src/hactoolnet/ProcessNca.cs b/src/hactoolnet/ProcessNca.cs
index 5cdf3d0b..c07f04c4 100644
--- a/src/hactoolnet/ProcessNca.cs
+++ b/src/hactoolnet/ProcessNca.cs
@@ -1,5 +1,4 @@
 using System.IO;
-using System.Linq;
 using System.Text;
 using LibHac;
 using LibHac.IO;
@@ -46,13 +45,13 @@ namespace hactoolnet
                         nca.ExtractSection(i, ctx.Options.SectionOutDir[i], ctx.Options.IntegrityLevel, ctx.Logger);
                     }
 
-                    if (ctx.Options.Validate && nca.Sections[i] != null)
+                    if (ctx.Options.Validate && nca.SectionExists(i))
                     {
                         nca.VerifySection(i, ctx.Logger);
                     }
                 }
 
-                if (ctx.Options.ListRomFs && nca.Sections[1] != null)
+                if (ctx.Options.ListRomFs && nca.CanOpenSection(NcaSectionType.Data))
                 {
                     IFileSystem romfs = nca.OpenFileSystem(NcaSectionType.Data, ctx.Options.IntegrityLevel);
 
@@ -64,35 +63,26 @@ namespace hactoolnet
 
                 if (ctx.Options.RomfsOutDir != null || ctx.Options.RomfsOut != null || ctx.Options.ReadBench)
                 {
-                    NcaSection section = nca.Sections.FirstOrDefault(x => x?.Type == SectionType.Romfs || x?.Type == SectionType.Bktr);
-
-                    if (section == null)
+                    if (!nca.SectionExists(NcaSectionType.Data))
                     {
                         ctx.Logger.LogMessage("NCA has no RomFS section");
                         return;
                     }
 
-                    if (section.Type == SectionType.Bktr && ctx.Options.BaseNca == null)
-                    {
-                        ctx.Logger.LogMessage("Cannot save BKTR section without base RomFS");
-                        return;
-                    }
-
                     if (ctx.Options.RomfsOut != null)
                     {
-                        nca.ExportSection(section.SectionNum, ctx.Options.RomfsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
+                        nca.ExportSection(NcaSectionType.Data, ctx.Options.RomfsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
                     }
 
                     if (ctx.Options.RomfsOutDir != null)
                     {
-                        IFileSystem romfs = nca.OpenFileSystem(section.SectionNum, ctx.Options.IntegrityLevel);
-                        romfs.Extract(ctx.Options.RomfsOutDir, ctx.Logger);
+                        nca.ExtractSection(NcaSectionType.Data, ctx.Options.RomfsOutDir, ctx.Options.IntegrityLevel, ctx.Logger);
                     }
 
                     if (ctx.Options.ReadBench)
                     {
                         long bytesToRead = 1024L * 1024 * 1024 * 5;
-                        IStorage storage = nca.OpenStorage(section.SectionNum, ctx.Options.IntegrityLevel);
+                        IStorage storage = nca.OpenStorage(NcaSectionType.Data, ctx.Options.IntegrityLevel);
                         var dest = new NullStorage(storage.Length);
 
                         int iterations = (int)(bytesToRead / storage.Length) + 1;
@@ -119,9 +109,7 @@ namespace hactoolnet
                         return;
                     }
 
-                    NcaSection section = nca.Sections[(int)ProgramPartitionType.Code];
-
-                    if (section == null)
+                    if (!nca.SectionExists(NcaSectionType.Code))
                     {
                         ctx.Logger.LogMessage("Could not find an ExeFS section");
                         return;
@@ -129,13 +117,12 @@ namespace hactoolnet
 
                     if (ctx.Options.ExefsOut != null)
                     {
-                        nca.ExportSection(section.SectionNum, ctx.Options.ExefsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
+                        nca.ExportSection(NcaSectionType.Code, ctx.Options.ExefsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
                     }
 
                     if (ctx.Options.ExefsOutDir != null)
                     {
-                        IFileSystem pfs = nca.OpenFileSystem(section.SectionNum, ctx.Options.IntegrityLevel);
-                        pfs.Extract(ctx.Options.ExefsOutDir, ctx.Logger);
+                        nca.ExtractSection(NcaSectionType.Code, ctx.Options.ExefsOutDir, ctx.Options.IntegrityLevel, ctx.Logger);
                     }
                 }
 
@@ -199,7 +186,7 @@ namespace hactoolnet
                     NcaSection sect = nca.Sections[i];
                     if (sect == null) continue;
 
-                    bool isExefs = nca.Header.ContentType == ContentType.Program && i == (int)ProgramPartitionType.Code;
+                    bool isExefs = nca.Header.ContentType == ContentType.Program && i == 0;
 
                     sb.AppendLine($"    Section {i}:");
                     PrintItem(sb, colLen, "        Offset:", $"0x{sect.Offset:x12}");
diff --git a/src/hactoolnet/ProcessRomfs.cs b/src/hactoolnet/ProcessRomfs.cs
index 99b87e6a..09220669 100644
--- a/src/hactoolnet/ProcessRomfs.cs
+++ b/src/hactoolnet/ProcessRomfs.cs
@@ -10,13 +10,14 @@ namespace hactoolnet
         {
             using (var file = new LocalStorage(ctx.Options.InFile, FileAccess.Read))
             {
-                var romfs = new RomFsFileSystem(file);
-                Process(ctx, romfs);
+                Process(ctx, file);
             }
         }
 
-        public static void Process(Context ctx, RomFsFileSystem romfs)
+        public static void Process(Context ctx, IStorage romfsStorage)
         {
+            var romfs = new RomFsFileSystem(romfsStorage);
+
             if (ctx.Options.ListRomFs)
             {
                 foreach (DirectoryEntry entry in romfs.EnumerateEntries())
@@ -29,7 +30,6 @@ namespace hactoolnet
             {
                 using (var outFile = new FileStream(ctx.Options.RomfsOut, FileMode.Create, FileAccess.ReadWrite))
                 {
-                    IStorage romfsStorage = romfs.GetBaseStorage();
                     romfsStorage.CopyToStream(outFile, romfsStorage.Length, ctx.Logger);
                 }
             }
diff --git a/src/hactoolnet/ProcessSwitchFs.cs b/src/hactoolnet/ProcessSwitchFs.cs
index 85528af4..ef8100ef 100644
--- a/src/hactoolnet/ProcessSwitchFs.cs
+++ b/src/hactoolnet/ProcessSwitchFs.cs
@@ -69,9 +69,7 @@ namespace hactoolnet
                     return;
                 }
 
-                NcaSection section = title.MainNca.Sections[(int)ProgramPartitionType.Code];
-
-                if (section == null)
+                if (!title.MainNca.SectionExists(NcaSectionType.Code))
                 {
                     ctx.Logger.LogMessage($"Main NCA for title {id:X16} has no ExeFS section");
                     return;
@@ -79,12 +77,12 @@ namespace hactoolnet
 
                 if (ctx.Options.ExefsOutDir != null)
                 {
-                    title.MainNca.ExtractSection(section.SectionNum, ctx.Options.ExefsOutDir, ctx.Options.IntegrityLevel, ctx.Logger);
+                    title.MainNca.ExtractSection(NcaSectionType.Code, ctx.Options.ExefsOutDir, ctx.Options.IntegrityLevel, ctx.Logger);
                 }
 
                 if (ctx.Options.ExefsOut != null)
                 {
-                    title.MainNca.ExportSection(section.SectionNum, ctx.Options.ExefsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
+                    title.MainNca.ExportSection(NcaSectionType.Code, ctx.Options.ExefsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
                 }
             }
 
@@ -109,24 +107,13 @@ namespace hactoolnet
                     return;
                 }
 
-                NcaSection section = title.MainNca.Sections.FirstOrDefault(x => x?.Type == SectionType.Romfs || x?.Type == SectionType.Bktr);
-
-                if (section == null)
+                if (!title.MainNca.SectionExists(NcaSectionType.Data))
                 {
                     ctx.Logger.LogMessage($"Main NCA for title {id:X16} has no RomFS section");
                     return;
                 }
 
-                if (ctx.Options.RomfsOutDir != null)
-                {
-                    IFileSystem romfs = title.MainNca.OpenFileSystem(section.SectionNum, ctx.Options.IntegrityLevel);
-                    romfs.Extract(ctx.Options.RomfsOutDir, ctx.Logger);
-                }
-
-                if (ctx.Options.RomfsOut != null)
-                {
-                    title.MainNca.ExportSection(section.SectionNum, ctx.Options.RomfsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
-                }
+                ProcessRomfs.Process(ctx, title.MainNca.OpenStorage(NcaSectionType.Data, ctx.Options.IntegrityLevel, false));
             }
 
             if (ctx.Options.OutDir != null)
diff --git a/src/hactoolnet/ProcessXci.cs b/src/hactoolnet/ProcessXci.cs
index cca5bbce..b6928c4f 100644
--- a/src/hactoolnet/ProcessXci.cs
+++ b/src/hactoolnet/ProcessXci.cs
@@ -72,9 +72,7 @@ namespace hactoolnet
                         return;
                     }
 
-                    NcaSection exefsSection = mainNca.Sections[(int)ProgramPartitionType.Code];
-
-                    if (exefsSection == null)
+                    if (!mainNca.SectionExists(NcaSectionType.Code))
                     {
                         ctx.Logger.LogMessage("NCA has no ExeFS section");
                         return;
@@ -82,12 +80,12 @@ namespace hactoolnet
 
                     if (ctx.Options.ExefsOutDir != null)
                     {
-                        mainNca.ExtractSection(exefsSection.SectionNum, ctx.Options.ExefsOutDir, ctx.Options.IntegrityLevel, ctx.Logger);
+                        mainNca.ExtractSection(NcaSectionType.Code, ctx.Options.ExefsOutDir, ctx.Options.IntegrityLevel, ctx.Logger);
                     }
 
                     if (ctx.Options.ExefsOut != null)
                     {
-                        mainNca.ExportSection(exefsSection.SectionNum, ctx.Options.ExefsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
+                        mainNca.ExportSection(NcaSectionType.Code, ctx.Options.ExefsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
                     }
                 }
 
@@ -101,24 +99,13 @@ namespace hactoolnet
                         return;
                     }
 
-                    NcaSection romfsSection = mainNca.Sections.FirstOrDefault(x => x.Type == SectionType.Romfs);
-
-                    if (romfsSection == null)
+                    if (!mainNca.SectionExists(NcaSectionType.Data))
                     {
                         ctx.Logger.LogMessage("NCA has no RomFS section");
                         return;
                     }
 
-                    if (ctx.Options.RomfsOutDir != null)
-                    {
-                        IFileSystem romfs = mainNca.OpenFileSystem(romfsSection.SectionNum, ctx.Options.IntegrityLevel);
-                        romfs.Extract(ctx.Options.RomfsOutDir, ctx.Logger);
-                    }
-
-                    if (ctx.Options.RomfsOut != null)
-                    {
-                        mainNca.ExportSection(romfsSection.SectionNum, ctx.Options.RomfsOut, ctx.Options.Raw, ctx.Options.IntegrityLevel, ctx.Logger);
-                    }
+                    ProcessRomfs.Process(ctx, mainNca.OpenStorage(NcaSectionType.Data, ctx.Options.IntegrityLevel, false));
                 }
             }
         }