First attempt at implementing BKTR

hactoolnet/CliParser.cs

@@ -30,6 +30,7 @@ namespace hactoolnet
             new CliOption("outdir", 1, (o, a) => o.OutDir = a[0]),
             new CliOption("sdseed", 1, (o, a) => o.SdSeed = a[0]),
             new CliOption("sdpath", 1, (o, a) => o.SdPath = a[0]),
+            new CliOption("basenca", 1, (o, a) => o.BaseNca = a[0]),
             new CliOption("listapps", 0, (o, a) => o.ListApps = true),
             new CliOption("listtitles", 0, (o, a) => o.ListTitles = true),
             new CliOption("listromfs", 0, (o, a) => o.ListRomFs = true),

hactoolnet/Options.cs

@@ -20,6 +20,7 @@ namespace hactoolnet
         public string OutDir;
         public string SdSeed;
         public string SdPath;
+        public string BaseNca;
         public bool ListApps;
         public bool ListTitles;
         public bool ListRomFs;

hactoolnet/Program.cs

@@ -52,6 +52,13 @@ namespace hactoolnet
             {
                 var nca = new Nca(ctx.Keyset, file, false);
 
+                if (ctx.Options.BaseNca != null)
+                {
+                    var baseFile = new FileStream(ctx.Options.BaseNca, FileMode.Open, FileAccess.Read);
+                    var baseNca = new Nca(ctx.Keyset, baseFile, false);
+                    nca.SetBaseNca(baseNca);
+                }
+
                 for (int i = 0; i < 3; i++)
                 {
                     if (ctx.Options.SectionOut[i] != null)
@@ -80,6 +87,36 @@ namespace hactoolnet
                     }
                 }
 
+                if (ctx.Options.RomfsOutDir != null)
+                {
+                    NcaSection section = nca.Sections.FirstOrDefault(x => x.Type == SectionType.Romfs || x.Type == SectionType.Bktr);
+
+                    if (section == null)
+                    {
+                        ctx.Logger.LogMessage("NCA has no RomFS section");
+                        return;
+                    }
+
+                    if (section.Type == SectionType.Bktr)
+                    {
+                        if (ctx.Options.BaseNca == null)
+                        {
+                            ctx.Logger.LogMessage("Cannot save BKTR section without base RomFS");
+                            return;
+                        }
+
+                        var bktr = nca.OpenSection(1, false);
+                        var romfs = new Romfs(bktr);
+                        romfs.Extract(ctx.Options.RomfsOutDir, ctx.Logger);
+
+                    }
+                    else
+                    {
+                        var romfs = new Romfs(nca.OpenSection(section.SectionNum, false));
+                        romfs.Extract(ctx.Options.RomfsOutDir, ctx.Logger);
+                    }
+                }
+
                 ctx.Logger.LogMessage(nca.Dump());
             }
         }

libhac/Aes128CounterMode.cs

@@ -66,6 +66,14 @@ namespace libhac
             }
         }
 
+        public void UpdateCounterSubsection(uint value)
+        {
+            _counter[7] = (byte) value;
+            _counter[6] = (byte) (value >> 8);
+            _counter[5] = (byte) (value >> 16);
+            _counter[4] = (byte) (value >> 24);
+        }
+
         private void EncryptCounterThenIncrement()
         {
             _counterEncryptor.TransformBlock(_counter, 0, _counter.Length, _counterEnc, 0);

libhac/AesCtrStream.cs

@@ -45,7 +45,7 @@ namespace libhac
         private readonly long _counterOffset;
         private readonly byte[] _tempBuffer;
         private readonly Aes _aes;
-        private CounterModeCryptoTransform _decryptor;
+        protected CounterModeCryptoTransform Decryptor;
 
         /// <summary>
         /// Creates a new stream
@@ -83,7 +83,7 @@ namespace libhac
             _aes.Key = key;
             _aes.Mode = CipherMode.ECB;
             _aes.Padding = PaddingMode.None;
-            _decryptor = CreateDecryptor();
+            Decryptor = CreateDecryptor();
 
         }
 
@@ -101,7 +101,7 @@ namespace libhac
         protected override void Dispose(bool disposing)
         {
             base.Dispose(disposing);
-            _decryptor?.Dispose();
+            Decryptor?.Dispose();
         }
 
         public override void Flush()
@@ -141,7 +141,7 @@ namespace libhac
             set
             {
                 base.Position = value;
-                _decryptor.UpdateCounter(_counterOffset + base.Position);
+                Decryptor.UpdateCounter(_counterOffset + base.Position);
             }
         }
 
@@ -162,11 +162,11 @@ namespace libhac
             if (ret == 0)
                 return 0;
 
-            if (_decryptor == null)
-                _decryptor = CreateDecryptor();
+            if (Decryptor == null)
+                Decryptor = CreateDecryptor();
 
             //decrypt the sector
-            var retV = _decryptor.TransformBlock(_tempBuffer, 0, buffer, offset);
+            var retV = Decryptor.TransformBlock(_tempBuffer, 0, buffer, offset);
 
             //Console.WriteLine("Decrypting sector {0} == {1} bytes", currentSector, retV);
 

libhac/Bktr.cs

@@ -0,0 +1,172 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Text;
+
+namespace libhac
+{
+    public class Bktr : Stream
+    {
+        private long _position;
+        public RelocationBlock RelocationBlock { get; }
+        private List<RelocationEntry> RelocationEntries { get; } = new List<RelocationEntry>();
+        private List<long> RelocationOffsets { get; }
+
+        private Stream Patch { get; }
+        private Stream Base { get; }
+        private RelocationEntry CurrentEntry { get; set; }
+
+        public Bktr(Stream patchRomfs, Stream baseRomfs, NcaSection section)
+        {
+            if (section.Type != SectionType.Bktr) throw new ArgumentException("Section is not of type BKTR");
+
+            Patch = patchRomfs;
+            Base = baseRomfs;
+            IvfcLevelHeader level5 = section.Header.Bktr.IvfcHeader.LevelHeaders[5];
+            Length = level5.LogicalOffset + level5.HashDataSize;
+
+            using (var reader = new BinaryReader(patchRomfs, Encoding.Default, true))
+            {
+                patchRomfs.Position = section.Header.Bktr.RelocationHeader.Offset;
+                RelocationBlock = new RelocationBlock(reader);
+            }
+
+            foreach (RelocationBucket bucket in RelocationBlock.Buckets)
+            {
+                RelocationEntries.AddRange(bucket.Entries);
+            }
+
+            for (int i = 0; i < RelocationEntries.Count - 1; i++)
+            {
+                RelocationEntries[i].Next = RelocationEntries[i + 1];
+                RelocationEntries[i].VirtOffsetEnd = RelocationEntries[i + 1].VirtOffset;
+            }
+
+            RelocationEntries[RelocationEntries.Count - 1].VirtOffsetEnd = level5.LogicalOffset + level5.HashDataSize;
+            RelocationOffsets = RelocationEntries.Select(x => x.VirtOffset).ToList();
+
+            CurrentEntry = GetRelocationEntry(0);
+            UpdateSourceStreamPositions();
+        }
+
+        private RelocationEntry GetRelocationEntry(long offset)
+        {
+            var index = RelocationOffsets.BinarySearch(offset);
+            if (index < 0) index = ~index - 1;
+            return RelocationEntries[index];
+        }
+
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            long remaining = Length - Position;
+            if (remaining <= 0) return 0;
+            if (remaining < count) count = (int)remaining;
+
+            var toOutput = count;
+            int pos = 0;
+
+            while (toOutput > 0)
+            {
+                var remainInEntry = CurrentEntry.VirtOffsetEnd - Position;
+                int toRead = (int)Math.Min(toOutput, remainInEntry);
+                ReadCurrent(buffer, pos, toRead);
+                pos += toRead;
+                toOutput -= toRead;
+            }
+
+            return count;
+        }
+
+        private void ReadCurrent(byte[] buffer, int offset, int count)
+        {
+            if (CurrentEntry.IsPatch)
+            {
+                Patch.Read(buffer, offset, count);
+            }
+            else
+            {
+                Base.Read(buffer, offset, count);
+            }
+
+            Position += count;
+        }
+
+        private void UpdateSourceStreamPositions()
+        {
+            // At end of virtual stream
+            if (CurrentEntry == null) return;
+
+            var entryOffset = Position - CurrentEntry.VirtOffset;
+
+            if (CurrentEntry.IsPatch)
+            {
+                Patch.Position = CurrentEntry.PhysOffset + entryOffset;
+            }
+            else
+            {
+                Base.Position = CurrentEntry.PhysOffset + entryOffset;
+            }
+        }
+
+        public override long Position
+        {
+            get => _position;
+            set
+            {
+                if (value > Length) throw new IndexOutOfRangeException();
+
+                // Avoid doing a search when reading sequentially
+                if (CurrentEntry != null && value == CurrentEntry.VirtOffsetEnd)
+                {
+                    CurrentEntry = CurrentEntry.Next;
+                }
+                else if (CurrentEntry == null || value < CurrentEntry.VirtOffset || value > CurrentEntry.VirtOffsetEnd)
+                {
+                    CurrentEntry = GetRelocationEntry(value);
+                }
+
+                _position = value;
+                UpdateSourceStreamPositions();
+            }
+        }
+
+        public override long Seek(long offset, SeekOrigin origin)
+        {
+            switch (origin)
+            {
+                case SeekOrigin.Begin:
+                    Position = offset;
+                    break;
+                case SeekOrigin.Current:
+                    Position += offset;
+                    break;
+                case SeekOrigin.End:
+                    Position = Length - offset;
+                    break;
+            }
+
+            return Position;
+        }
+
+        public override void SetLength(long value)
+        {
+            throw new NotSupportedException();
+        }
+
+        public override void Write(byte[] buffer, int offset, int count)
+        {
+            throw new NotSupportedException();
+        }
+
+        public override void Flush()
+        {
+            throw new NotSupportedException();
+        }
+
+        public override bool CanRead => true;
+        public override bool CanWrite => false;
+        public override long Length { get; }
+        public override bool CanSeek => true;
+    }
+}

libhac/BktrCryptoStream.cs

@@ -0,0 +1,93 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using libhac.XTSSharp;
+
+namespace libhac
+{
+    public class BktrCryptoStream : AesCtrStream
+    {
+        public SubsectionBlock SubsectionBlock { get; }
+        private List<SubsectionEntry> SubsectionEntries { get; } = new List<SubsectionEntry>();
+        private List<long> SubsectionOffsets { get; }
+        private SubsectionEntry CurrentEntry { get; set; }
+
+        public BktrCryptoStream(Stream baseStream, byte[] key, long offset, long length, long counterOffset, byte[] ctrHi, NcaSection section)
+            : base(baseStream, key, offset, length, counterOffset, ctrHi)
+        {
+            if (section.Type != SectionType.Bktr) throw new ArgumentException("Section is not of type BKTR");
+
+            var bktr = section.Header.Bktr;
+            var header = bktr.SubsectionHeader;
+            byte[] subsectionBytes;
+            using (var streamDec = new RandomAccessSectorStream(new AesCtrStream(baseStream, key, offset, length, counterOffset, ctrHi)))
+            {
+                streamDec.Position = header.Offset;
+                subsectionBytes = new byte[header.Size];
+                streamDec.Read(subsectionBytes, 0, subsectionBytes.Length);
+            }
+
+            using (var reader = new BinaryReader(new MemoryStream(subsectionBytes)))
+            {
+                SubsectionBlock = new SubsectionBlock(reader);
+            }
+
+            foreach (var bucket in SubsectionBlock.Buckets)
+            {
+                SubsectionEntries.AddRange(bucket.Entries);
+            }
+
+            // Add a subsection for the BKTR headers to make things easier
+            var headerSubsection = new SubsectionEntry
+            {
+                Offset = bktr.RelocationHeader.Offset,
+                Counter = (uint)(ctrHi[4] << 24 | ctrHi[5] << 16 | ctrHi[6] << 8 | ctrHi[7]),
+                OffsetEnd = long.MaxValue
+            };
+            SubsectionEntries.Add(headerSubsection);
+
+            for (int i = 0; i < SubsectionEntries.Count - 1; i++)
+            {
+                SubsectionEntries[i].Next = SubsectionEntries[i + 1];
+                SubsectionEntries[i].OffsetEnd = SubsectionEntries[i + 1].Offset;
+            }
+
+            SubsectionOffsets = SubsectionEntries.Select(x => x.Offset).ToList();
+
+            CurrentEntry = GetSubsectionEntry(0);
+            Decryptor.UpdateCounterSubsection(CurrentEntry.Counter);
+            baseStream.Position = offset;
+        }
+
+        public override long Position
+        {
+            get => base.Position;
+            set
+            {
+                base.Position = value;
+                CurrentEntry = GetSubsectionEntry(value);
+                Decryptor.UpdateCounterSubsection(CurrentEntry.Counter);
+            }
+        }
+
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            var ret = base.Read(buffer, offset, count);
+            if (Position >= CurrentEntry.OffsetEnd)
+            {
+                CurrentEntry = CurrentEntry.Next;
+                Decryptor.UpdateCounterSubsection(CurrentEntry.Counter);
+            }
+
+            return ret;
+        }
+
+        private SubsectionEntry GetSubsectionEntry(long offset)
+        {
+            var index = SubsectionOffsets.BinarySearch(offset);
+            if (index < 0) index = ~index - 1;
+            return SubsectionEntries[index];
+        }
+    }
+}

libhac/BktrStructs.cs

@@ -0,0 +1,152 @@
+using System.IO;
+
+namespace libhac
+{
+    public class RelocationBlock
+    {
+        public uint Field0;
+        public int BucketCount;
+        public long Size;
+        public long[] BaseOffsets;
+        public RelocationBucket[] Buckets;
+
+        public RelocationBlock(BinaryReader reader)
+        {
+            var start = reader.BaseStream.Position;
+
+            Field0 = reader.ReadUInt32();
+            BucketCount = reader.ReadInt32();
+            Size = reader.ReadInt64();
+            BaseOffsets = new long[BucketCount];
+            Buckets = new RelocationBucket[BucketCount];
+
+            for (int i = 0; i < BucketCount; i++)
+            {
+                BaseOffsets[i] = reader.ReadInt64();
+            }
+
+            reader.BaseStream.Position = start + 0x4000;
+
+            for (int i = 0; i < BucketCount; i++)
+            {
+                Buckets[i] = new RelocationBucket(reader);
+            }
+        }
+    }
+
+    public class RelocationBucket
+    {
+        public int BucketNum;
+        public int EntryCount;
+        public long VirtualOffsetEnd;
+        public RelocationEntry[] Entries;
+
+        public RelocationBucket(BinaryReader reader)
+        {
+            var start = reader.BaseStream.Position;
+
+            BucketNum = reader.ReadInt32();
+            EntryCount = reader.ReadInt32();
+            VirtualOffsetEnd = reader.ReadInt64();
+            Entries = new RelocationEntry[EntryCount];
+
+            for (int i = 0; i < EntryCount; i++)
+            {
+                Entries[i] = new RelocationEntry(reader);
+            }
+
+            reader.BaseStream.Position = start + 0x4000;
+        }
+    }
+
+    public class RelocationEntry
+    {
+        public long VirtOffset;
+        public long VirtOffsetEnd;
+        public long PhysOffset;
+        public bool IsPatch;
+        public RelocationEntry Next;
+
+        public RelocationEntry(BinaryReader reader)
+        {
+            VirtOffset = reader.ReadInt64();
+            PhysOffset = reader.ReadInt64();
+            IsPatch = reader.ReadInt32() != 0;
+        }
+    }
+
+    public class SubsectionBlock
+    {
+        public uint Field0;
+        public int BucketCount;
+        public long Size;
+        public long[] BaseOffsets;
+        public SubsectionBucket[] Buckets;
+
+        public SubsectionBlock(BinaryReader reader)
+        {
+            var start = reader.BaseStream.Position;
+
+            Field0 = reader.ReadUInt32();
+            BucketCount = reader.ReadInt32();
+            Size = reader.ReadInt64();
+            BaseOffsets = new long[BucketCount];
+            Buckets = new SubsectionBucket[BucketCount];
+
+            for (int i = 0; i < BucketCount; i++)
+            {
+                BaseOffsets[i] = reader.ReadInt64();
+            }
+
+            reader.BaseStream.Position = start + 0x4000;
+
+            for (int i = 0; i < BucketCount; i++)
+            {
+                Buckets[i] = new SubsectionBucket(reader);
+            }
+        }
+    }
+
+    public class SubsectionBucket
+    {
+        public int BucketNum;
+        public int EntryCount;
+        public long VirtualOffsetEnd;
+        public SubsectionEntry[] Entries;
+        public SubsectionBucket(BinaryReader reader)
+        {
+            var start = reader.BaseStream.Position;
+
+            BucketNum = reader.ReadInt32();
+            EntryCount = reader.ReadInt32();
+            VirtualOffsetEnd = reader.ReadInt64();
+            Entries = new SubsectionEntry[EntryCount];
+
+            for (int i = 0; i < EntryCount; i++)
+            {
+                Entries[i] = new SubsectionEntry(reader);
+            }
+
+            reader.BaseStream.Position = start + 0x4000;
+        }
+    }
+
+    public class SubsectionEntry
+    {
+        public long Offset;
+        public uint Field8;
+        public uint Counter;
+
+        public SubsectionEntry Next;
+        public long OffsetEnd;
+
+        public SubsectionEntry() { }
+
+        public SubsectionEntry(BinaryReader reader)
+        {
+            Offset = reader.ReadInt64();
+            Field8 = reader.ReadUInt32();
+            Counter = reader.ReadUInt32();
+        }
+    }
+}

libhac/Nca.cs

@@ -1,5 +1,6 @@
 using System;
 using System.IO;
+using System.Linq;
 using System.Security.Cryptography;
 using System.Text;
 using libhac.XTSSharp;
@@ -18,6 +19,7 @@ namespace libhac
         public byte[] TitleKeyDec { get; } = new byte[0x10];
         public Stream Stream { get; private set; }
         private bool KeepOpen { get; }
+        private Nca BaseNca { get; set; }
 
         public NcaSection[] Sections { get; } = new NcaSection[4];
 
@@ -77,9 +79,6 @@ namespace libhac
                         size = sect.Header.Romfs.IvfcHeader.LevelHeaders[Romfs.IvfcMaxLevel - 1].HashDataSize;
                         break;
                     case SectionType.Bktr:
-                        offset = sect.Offset + sect.Header.Bktr.IvfcHeader.LevelHeaders[Romfs.IvfcMaxLevel - 1]
-                                     .LogicalOffset;
-                        size = sect.Header.Bktr.IvfcHeader.LevelHeaders[Romfs.IvfcMaxLevel - 1].HashDataSize;
                         break;
                     default:
                         throw new ArgumentOutOfRangeException();
@@ -97,7 +96,27 @@ namespace libhac
                 case SectionCryptType.CTR:
                     return new RandomAccessSectorStream(new AesCtrStream(Stream, DecryptedKeys[2], offset, size, offset, sect.Header.Ctr), false);
                 case SectionCryptType.BKTR:
-                    return new RandomAccessSectorStream(new AesCtrStream(Stream, DecryptedKeys[2], offset, size, offset, sect.Header.Ctr), false);
+                    var patchStream = new RandomAccessSectorStream(
+                        new BktrCryptoStream(Stream, DecryptedKeys[2], offset, size, offset, sect.Header.Ctr, sect),
+                        false);
+                    if (BaseNca == null)
+                    {
+                        return patchStream;
+                    }
+                    else
+                    {
+                        var dataLevel = sect.Header.Bktr.IvfcHeader.LevelHeaders[Romfs.IvfcMaxLevel - 1];
+
+                        var baseSect = BaseNca.Sections.FirstOrDefault(x => x.Type == SectionType.Romfs);
+                        if (baseSect == null) throw new InvalidDataException("Base NCA has no RomFS section");
+
+                        var baseStream = BaseNca.OpenSection(baseSect.SectionNum, true);
+                        var virtStreamRaw = new Bktr(patchStream, baseStream, sect);
+
+                        if (raw) return virtStreamRaw;
+                        var virtStream = new SubStream(virtStreamRaw, dataLevel.LogicalOffset, dataLevel.HashDataSize);
+                        return virtStream;
+                    }
                 default:
                     throw new ArgumentOutOfRangeException();
             }
@@ -105,6 +124,8 @@ namespace libhac
             return new SubStream(Stream, offset, size);
         }
 
+        public void SetBaseNca(Nca baseNca) => BaseNca = baseNca;
+
         private void DecryptHeader(Keyset keyset, Stream stream)
         {
             byte[] headerBytes = new byte[0xC00];

libhac/NcaStructs.cs

@@ -204,8 +204,8 @@ namespace libhac
 
     public class BktrHeader
     {
-        public ulong Offset;
-        public ulong Size;
+        public long Offset;
+        public long Size;
         public uint Magic;
         public uint Field14;
         public uint NumEntries;
@@ -213,8 +213,8 @@ namespace libhac
 
         public BktrHeader(BinaryReader reader)
         {
-            Offset = reader.ReadUInt64();
-            Size = reader.ReadUInt64();
+            Offset = reader.ReadInt64();
+            Size = reader.ReadInt64();
             Magic = reader.ReadUInt32();
             Field14 = reader.ReadUInt32();
             NumEntries = reader.ReadUInt32();

libhac/Substream.cs

@@ -20,6 +20,7 @@ namespace libhac
 
             baseStream.Seek(offset, SeekOrigin.Begin);
         }
+
         public override int Read(byte[] buffer, int offset, int count)
         {
             long remaining = Length - Position;