TinfoilVibeServer/TinfoilVibeServer/Services/NSPExtractor.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using LibHac.Common;
using LibHac.Fs;
using LibHac.Fs.Fsa;
using LibHac.FsSystem;
using LibHac.Tools.FsSystem;
using LibHac.Tools.FsSystem.NcaUtils;
using LibHac.Common.Keys;
using LibHac.Ncm;
using LibHac.Tools.Fs;
using LibHac.Tools.Ncm;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using TinfoilVibeServer.Models;
using TinfoilVibeServer.Utilities;
using NcaHeader = LibHac.Tools.FsSystem.NcaUtils.NcaHeader;
using Path = System.IO.Path;

namespace TinfoilVibeServer.Services
{
    public interface INSPExtractor
    {
        /// <summary>
        /// Public convenience wrapper that opens the file on disk.
        /// </summary>
        NcaMetadataWithHash? ExtractFromFile(string filePath);

        /// <summary>
        /// Core implementation – works on any seekable stream that contains a full NSP/XCI container.
        /// </summary>
        NcaMetadataWithHash? ExtractFromStream(Stream stream);

        string ExtractHashFromStream(Stream nspStream);
    }

    /// <summary>
    /// Extracts the TitleId, version, type *and* the SHA‑256 of the first NCA stream.
    /// </summary>
    public sealed class NSPExtractor : INSPExtractor
    {
        private KeySet? _keySet;

        public KeySet? KeySet
        {
            get
            {
                if (_keySet != null) return _keySet;
                if (_options.CurrentValue.KeyFile == null) return null;
                var dataRoot = _environment.ContentRootPath ?? "/app/config";
                if (Path.IsPathRooted(_options.CurrentValue.KeyFile))
                {
                    _keySet = ExternalKeyReader.ReadKeyFile(_options.CurrentValue.KeyFile);
                }
                else
                {
                    _keySet = ExternalKeyReader.ReadKeyFile(Path.Combine(dataRoot, "config", _options.CurrentValue.KeyFile));
                }

                return _keySet;
            }
        }

        private readonly IOptionsMonitor<NSPExtractorOptions> _options;
        private readonly ILogger<INSPExtractor> _logger;
        private readonly IHostEnvironment _environment;

        public NSPExtractor(IOptionsMonitor<NSPExtractorOptions> options, ILogger<INSPExtractor> logger, IHostEnvironment environment)
        {
            _options = options;
            _options.OnChange(o =>
            {
                if (o.KeyFile == null)
                {
                    _logger?.LogInformation("No KeySet specified, skipping key validation");
                }

                if (!File.Exists(o.KeyFile))
                {
                    _logger?.LogWarning("KeySet file {KeyFile} does not exist", o.KeyFile);
                }
            });
            _logger = logger;
            _environment = environment;
        }

        /// <summary>
        /// Public convenience wrapper that opens the file on disk.
        /// </summary>
        public NcaMetadataWithHash? ExtractFromFile(string filePath)
        {
            using var stream = File.OpenRead(filePath);
            return ExtractFromStream(stream);
        }

        /// <summary>
        /// Core implementation – works on any seekable stream that contains a full NSP/XCI container.
        /// </summary>
        public NcaMetadataWithHash? ExtractFromStream(Stream stream)
        {
            if (KeySet == null) return null;

            if (!stream.CanSeek) return null;
            stream.Seek(0, SeekOrigin.Begin);

            _logger.LogInformation("Extracting NSP/XCI from stream (length={Length})", stream.Length);
            using var storage = new StreamStorage(stream, false);
            if (IsPfs0FileSystem(stream))
            {
                return ExtractNSPFromStream(storage);
            }

            if (IsXciFileSystem(stream))
            {
                var xci = new Xci(KeySet, storage);
                if (xci.HasPartition(XciPartitionType.Secure))
                {
                    _logger.LogInformation("Processing as XCI");
                    using var partition = xci.OpenPartition(XciPartitionType.Secure);
                    // Find the smallest file
                    var hashEntry = partition.EnumerateEntries("*", SearchOptions.Default)
                        .First(e => e.Type == DirectoryEntryType.File && e.Size < 10 * 1024 * 1024);
                    using var hashFileRef = new UniqueRef<IFile>();
                    partition.OpenFile(ref hashFileRef.Ref, hashEntry.FullPath.ToU8Span(), OpenMode.Read).ThrowIfFailure();
                    using var hashFile = hashFileRef.Release();
                    using var sha256 = SHA256.Create();
                    using var hashStream = hashFile.AsStream();
                    var hash = sha256.ComputeHash(hashStream);

                    List<DirectoryEntryEx> ncaEntries = partition
                        .EnumerateEntries("*.cnmt.nca", SearchOptions.RecurseSubdirectories)
                        .Where(e => e.Type == DirectoryEntryType.File)
                        .ToList();
                    foreach (var dirEntry in ncaEntries)
                    {
                        using var fileRef = new UniqueRef<IFile>();
                        partition.OpenFile(ref fileRef.Ref, dirEntry.FullPath.ToU8Span(), OpenMode.Read).ThrowIfFailure();
                        using var ncaFile = fileRef.Release();
                        using var ncaFileStorage = new FileStorage(ncaFile);

                        var nca = new Nca(KeySet, ncaFileStorage);
                       
                        if (nca.Header.ContentType != NcaContentType.Meta)
                            continue; // only the meta NCA contains title metadata

                        string titleId = nca.Header.TitleId.ToString("X16");
                        var (contentMetaType, applicationTitleId, titleVersion) = GetMetaData(nca);

                        _logger.LogInformation("Meta NCA found – TitleId={TitleId} ApplicationTitleId={ApplicationTitleId} Version={Version}", titleId, applicationTitleId.ToStrId(), titleVersion);
                        // XCI can never be a patch?
                        return new NcaMetadataWithHash(titleId, applicationTitleId.ToString("X16"), titleVersion.Major, ContentMetaType.Application, BitConverter.ToString(hash).Replace("-", "").ToLowerInvariant());
                    }
                }
            }

            return null; // no meta NCA found
        }

        private NcaMetadataWithHash? ExtractNSPFromStream(StreamStorage storage)
        {
            if (KeySet == null) return null;

            List<DirectoryEntryEx> ncaEntries;
            _logger.LogInformation("Processing as NSP");
            using var partition = new PartitionFileSystem();
            partition.Initialize(storage).ThrowIfFailure();
            // Find the first *.nca that contains the meta header
            var hashEntry = partition
                .EnumerateEntries("*", SearchOptions.Default)
                .First(e => e.Type == DirectoryEntryType.File && e.Size < 10 * 1024 * 1024);
            // Hash the *first* NCA stream – the stream we just opened
            using var hashRef = new UniqueRef<IFile>();
            partition.OpenFile(ref hashRef.Ref, hashEntry.FullPath.ToU8Span(), OpenMode.Read).ThrowIfFailure();
            using var hashFile = hashRef.Release();
            using var sha256 = SHA256.Create();
            using var hashStream = hashFile.AsStream();
            var hash = sha256.ComputeHash(hashStream);
            
            ncaEntries = partition
                .EnumerateEntries("*.nca", SearchOptions.RecurseSubdirectories)
                .Where(e => e.Type == DirectoryEntryType.File)
                .ToList();

            foreach (var dirEntry in ncaEntries)
            {
                using var fileRef = new UniqueRef<IFile>();
                partition.OpenFile(ref fileRef.Ref, dirEntry.FullPath.ToU8Span(), OpenMode.Read).ThrowIfFailure();
                using var ncaFile = fileRef.Release();
                using var ncaFileStorage = new FileStorage(ncaFile);

                var nca = new Nca(KeySet, ncaFileStorage);
                

                if (nca.Header.ContentType != NcaContentType.Meta)
                    continue; // only the meta NCA contains title metadata

                _logger.LogInformation("Meta NCA found – TitleId={TitleId} Version={Version}", nca.Header.TitleId, nca.Header.Version);
                string titleId = nca.Header.TitleId.ToString("X16");

                var (contentMetaType, applicationTitle, titleVersion) = GetMetaData(nca);
                if (contentMetaType != null)
                    return new NcaMetadataWithHash(titleId, applicationTitle.ToString("X16"), titleVersion.Minor, contentMetaType.Value, BitConverter.ToString(hash).Replace("-", "").ToLowerInvariant());
            }

            return null;
        }

        private static (ContentMetaType?, ulong, TitleVersion) GetMetaData(Nca nca)
        {
            if (nca.Header.ContentType != NcaContentType.Meta) return (null, 0, new TitleVersion(0, true));
            using var openFileSystem = nca.OpenFileSystem(0, IntegrityCheckLevel.ErrorOnInvalid);
            foreach (var entry in openFileSystem.EnumerateEntries("*.cnmt", SearchOptions.Default))
            {
                using var fileRef = new UniqueRef<IFile>();

                var result = openFileSystem.OpenFile(ref fileRef.Ref, entry.FullPath.ToU8Span(), OpenMode.Read);
                if (result.IsFailure()) continue;
                using var nacpFile = fileRef.Release();
                using var asStream = nacpFile.AsStream();

                var cnmt = new Cnmt(asStream);
                var applicationTitle = cnmt.ApplicationTitleId;
                return (cnmt.Type, applicationTitle, cnmt.TitleVersion);
            }

            return (null, 0, new TitleVersion(0, true));
        }

        /// <summary>
        /// Quick sanity check that the stream looks like a PFS0 file system.
        /// </summary>
        private bool IsPfs0FileSystem(Stream stream)
        {
            try
            {
                if (!stream.CanSeek) return false;
                stream.Seek(0, SeekOrigin.Begin);

                using var storage = new StreamStorage(stream, true);
                using var partition = new PartitionFileSystem();
                partition.Initialize(storage).ThrowIfFailure();
                _logger.LogInformation("PFS0 found");
                return true;
            }
            catch
            {
                // ignored
            }

            return false;
        }

        private bool IsXciFileSystem(Stream stream)
        {
            if (KeySet == null) return false;

            try
            {
                if (!stream.CanSeek) return false;
                stream.Seek(0, SeekOrigin.Begin);
                using var storage = new StreamStorage(stream, true);
                try
                {
                    var xciBlock = new Xci(KeySet, storage);
                    _logger.LogInformation("XCI found");
                    return xciBlock.HasPartition(XciPartitionType.Secure);
                }
                catch
                {
                    // ignored
                }

                return false;
            }
            catch (Exception e)
            {
                _logger.LogError("Failed to extract XCI: {Exception}", e.Message);
                return false;
            }
        }

        public string ExtractHashFromStream(Stream nspStream)
        {
            if (KeySet == null) return string.Empty;
            var seekableStream = !nspStream.CanSeek ? new SeekableBufferedStream(nspStream, nspStream.Length) : nspStream;

            var isNsp = IsPfs0FileSystem(seekableStream);
            var isXci = false;

            if (!isNsp)
            {
                isXci = IsXciFileSystem(seekableStream);
            }

            if (!isNsp && !isXci)
                return string.Empty;

            seekableStream.Seek(0, SeekOrigin.Begin);

            using var storage = new StreamStorage(seekableStream, true);
            if (isXci)
            {
                var xci = new Xci(KeySet, storage);
                using IFileSystem partition = xci.OpenPartition(XciPartitionType.Secure);
                var hashEntry = partition
                    .EnumerateEntries("*", SearchOptions.Default)
                    .First(e => e.Type == DirectoryEntryType.File && e.Size < 10*1024*1024);

                using var fileRef = new UniqueRef<IFile>();
                partition.OpenFile(ref fileRef.Ref, hashEntry.FullPath.ToU8Span(), OpenMode.Read).ThrowIfFailure();
                using var ncaFile = fileRef.Release();
                using var sha256 = SHA256.Create();
                using var ncaStream = ncaFile.AsStream();
                var hash = sha256.ComputeHash(ncaStream);
                var extractHashFromStream = BitConverter.ToString(hash).Replace("-", "").ToLowerInvariant();
                _logger.LogInformation("Computed first‑stream hash {Hash}", extractHashFromStream);
                return extractHashFromStream;
            }

            if (isNsp)
            {
                using var partition = new PartitionFileSystem();
                partition.Initialize(storage).ThrowIfFailure();
                var hashEntry = partition
                    .EnumerateEntries("*", SearchOptions.Default)
                    .First(e => e.Type == DirectoryEntryType.File && e.Size < 10 * 1024 * 1024);
                
                using var fileRef = new UniqueRef<IFile>();
                partition.OpenFile(ref fileRef.Ref, hashEntry.FullPath.ToU8Span(), OpenMode.Read).ThrowIfFailure();
                using var ncaFile = fileRef.Release();
                try
                {
                    using var ncaStream = ncaFile.AsStream();
                    using var sha256 = SHA256.Create();
                    var hash = sha256.ComputeHash(ncaStream);
                    var extractHashFromStream = BitConverter.ToString(hash).Replace("-", "").ToLowerInvariant();
                    _logger.LogInformation("Computed first‑stream hash {Hash}", extractHashFromStream);
                    return extractHashFromStream;
                }
                catch (Exception e)
                {
                    _logger.LogError("Failed to extract NSP: {Exception}", e.Message);
                }
            }

            return string.Empty;
        }
    }

    public class NSPExtractorOptions    
    {
        public string? KeyFile { get; set; }
    }
}