From 2844466bbbcaaa674a9fb12953b4dc240be0dbf5 Mon Sep 17 00:00:00 2001
From: Alex Barney <thealexbarney@gmail.com>
Date: Thu, 7 Feb 2019 20:32:57 -0600
Subject: [PATCH] Fix AesXtsFileSystem bugs and add the option to specify keys

---
 src/LibHac/IO/AesXtsDirectory.cs  |  8 +++++---
 src/LibHac/IO/AesXtsFile.cs       | 11 ++++++++++-
 src/LibHac/IO/AesXtsFileHeader.cs |  6 +++---
 src/LibHac/IO/AesXtsFileSystem.cs | 17 +++++++++++++++--
 4 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/src/LibHac/IO/AesXtsDirectory.cs b/src/LibHac/IO/AesXtsDirectory.cs
index a8ec03cd..2d5227c7 100644
--- a/src/LibHac/IO/AesXtsDirectory.cs
+++ b/src/LibHac/IO/AesXtsDirectory.cs
@@ -5,14 +5,16 @@ namespace LibHac.IO
 {
     public class AesXtsDirectory : IDirectory
     {
-        public IFileSystem ParentFileSystem { get; }
+        IFileSystem IDirectory.ParentFileSystem => ParentFileSystem;
+        public AesXtsFileSystem ParentFileSystem { get; }
+
         public string FullPath { get; }
         public OpenDirectoryMode Mode { get; }
 
         private IFileSystem BaseFileSystem { get; }
         private IDirectory BaseDirectory { get; }
 
-        public AesXtsDirectory(IFileSystem parentFs, IDirectory baseDir, OpenDirectoryMode mode)
+        public AesXtsDirectory(AesXtsFileSystem parentFs, IDirectory baseDir, OpenDirectoryMode mode)
         {
             ParentFileSystem = parentFs;
             BaseDirectory = baseDir;
@@ -66,7 +68,7 @@ namespace LibHac.IO
                     if (BitConverter.ToUInt32(buffer, 0) != 0x3058414E) return 0;
 
                     file.Read(buffer, 0x48);
-                    return BitConverter.ToInt32(buffer, 0);
+                    return BitConverter.ToInt64(buffer, 0);
                 }
             }
             catch (ArgumentOutOfRangeException)
diff --git a/src/LibHac/IO/AesXtsFile.cs b/src/LibHac/IO/AesXtsFile.cs
index c211d755..f054193c 100644
--- a/src/LibHac/IO/AesXtsFile.cs
+++ b/src/LibHac/IO/AesXtsFile.cs
@@ -31,10 +31,19 @@ namespace LibHac.IO
                 throw new ArgumentException("NAX0 key derivation failed.");
             }
 
-            IStorage encStorage = BaseFile.AsStorage().Slice(HeaderLength, Header.Size);
+            IStorage encStorage = BaseFile.AsStorage().Slice(HeaderLength, Util.AlignUp(Header.Size, 0x10));
             BaseStorage = new CachedStorage(new Aes128XtsStorage(encStorage, Header.DecryptedKey1, Header.DecryptedKey2, BlockSize, true), 4, true);
         }
 
+        public byte[] GetKey()
+        {
+            var key = new byte[0x20];
+            Array.Copy(Header.DecryptedKey1, 0, key, 0, 0x10);
+            Array.Copy(Header.DecryptedKey2, 0, key, 0x10, 0x10);
+
+            return key;
+        }
+
         public override int Read(Span<byte> destination, long offset)
         {
             int toRead = ValidateReadParamsAndGetSize(destination, offset);
diff --git a/src/LibHac/IO/AesXtsFileHeader.cs b/src/LibHac/IO/AesXtsFileHeader.cs
index b77767d4..3d72827b 100644
--- a/src/LibHac/IO/AesXtsFileHeader.cs
+++ b/src/LibHac/IO/AesXtsFileHeader.cs
@@ -37,10 +37,10 @@ namespace LibHac.IO
             }
         }
 
-        public AesXtsFileHeader(byte[] key1, byte[] key2, long fileSize, string path, byte[] kekSeed, byte[] verificationKey)
+        public AesXtsFileHeader(byte[] key, long fileSize, string path, byte[] kekSeed, byte[] verificationKey)
         {
-            Array.Copy(key1, DecryptedKey1, 0x10);
-            Array.Copy(key2, DecryptedKey2, 0x10);
+            Array.Copy(key, 0, DecryptedKey1, 0, 0x10);
+            Array.Copy(key, 0x10, DecryptedKey2, 0, 0x10);
             Magic = AesXtsFileMagic;
             Size = fileSize;
 
diff --git a/src/LibHac/IO/AesXtsFileSystem.cs b/src/LibHac/IO/AesXtsFileSystem.cs
index e26953f7..d736e5cc 100644
--- a/src/LibHac/IO/AesXtsFileSystem.cs
+++ b/src/LibHac/IO/AesXtsFileSystem.cs
@@ -33,10 +33,23 @@ namespace LibHac.IO
 
         public void CreateFile(string path, long size, CreateFileOptions options)
         {
-            long containerSize = AesXtsFile.HeaderLength + Util.AlignUp(size, 0x16);
+            CreateFile(path, size, options, new byte[0x20]);
+        }
+
+        /// <summary>
+        /// Creates a new <see cref="AesXtsFile"/> using the provided key.
+        /// </summary>
+        /// <param name="path">The full path of the file to create.</param>
+        /// <param name="size">The initial size of the created file.</param>
+        /// <param name="options">Flags to control how the file is created.
+        /// Should usually be <see cref="CreateFileOptions.None"/></param>
+        /// <param name="key">The 256-bit key containing a 128-bit data key followed by a 128-bit tweak key.</param>
+        public void CreateFile(string path, long size, CreateFileOptions options, byte[] key)
+        {
+            long containerSize = AesXtsFile.HeaderLength + Util.AlignUp(size, 0x10);
             BaseFileSystem.CreateFile(path, containerSize, options);
 
-            var header = new AesXtsFileHeader(new byte[0x10], new byte[0x10], size, path, KekSource, ValidationKey);
+            var header = new AesXtsFileHeader(key, size, path, KekSource, ValidationKey);
 
             using (IFile baseFile = BaseFileSystem.OpenFile(path, OpenMode.Write))
             {