From 18bb3d853127fa9ea3dfd9fe8ca3040f473735da Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Sat, 7 Jul 2018 15:45:06 -0500 Subject: [PATCH] First attempt at implementing BKTR --- hactoolnet/CliParser.cs | 1 + hactoolnet/Options.cs | 1 + hactoolnet/Program.cs | 37 ++++++++ libhac/Aes128CounterMode.cs | 8 ++ libhac/AesCtrStream.cs | 14 +-- libhac/Bktr.cs | 172 ++++++++++++++++++++++++++++++++++++ libhac/BktrCryptoStream.cs | 93 +++++++++++++++++++ libhac/BktrStructs.cs | 152 +++++++++++++++++++++++++++++++ libhac/Nca.cs | 29 +++++- libhac/NcaStructs.cs | 8 +- libhac/Substream.cs | 1 + 11 files changed, 501 insertions(+), 15 deletions(-) create mode 100644 libhac/Bktr.cs create mode 100644 libhac/BktrCryptoStream.cs create mode 100644 libhac/BktrStructs.cs diff --git a/hactoolnet/CliParser.cs b/hactoolnet/CliParser.cs index 9606e6b5..85ce897e 100644 --- a/hactoolnet/CliParser.cs +++ b/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), diff --git a/hactoolnet/Options.cs b/hactoolnet/Options.cs index a9103446..2627827d 100644 --- a/hactoolnet/Options.cs +++ b/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; diff --git a/hactoolnet/Program.cs b/hactoolnet/Program.cs index 89de547c..9c61366a 100644 --- a/hactoolnet/Program.cs +++ b/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()); } } diff --git a/libhac/Aes128CounterMode.cs b/libhac/Aes128CounterMode.cs index e909dffa..542db009 100644 --- a/libhac/Aes128CounterMode.cs +++ b/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); diff --git a/libhac/AesCtrStream.cs b/libhac/AesCtrStream.cs index 8165d8b8..e26e79b4 100644 --- a/libhac/AesCtrStream.cs +++ b/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; /// /// 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); diff --git a/libhac/Bktr.cs b/libhac/Bktr.cs new file mode 100644 index 00000000..442cb35a --- /dev/null +++ b/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 RelocationEntries { get; } = new List(); + private List 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; + } +} diff --git a/libhac/BktrCryptoStream.cs b/libhac/BktrCryptoStream.cs new file mode 100644 index 00000000..4ae44955 --- /dev/null +++ b/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 SubsectionEntries { get; } = new List(); + private List 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]; + } + } +} diff --git a/libhac/BktrStructs.cs b/libhac/BktrStructs.cs new file mode 100644 index 00000000..5df564d8 --- /dev/null +++ b/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(); + } + } +} diff --git a/libhac/Nca.cs b/libhac/Nca.cs index e10a6229..24e3af40 100644 --- a/libhac/Nca.cs +++ b/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]; diff --git a/libhac/NcaStructs.cs b/libhac/NcaStructs.cs index 7c7e08fa..978e109a 100644 --- a/libhac/NcaStructs.cs +++ b/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(); diff --git a/libhac/Substream.cs b/libhac/Substream.cs index 57afae9f..7dbeb586 100644 --- a/libhac/Substream.cs +++ b/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;