TinfoilVibeServer/TinfoilVibeServer/Services/ArchiveHandler.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using Microsoft.Extensions.Logging;
using SharpCompress.Archives;
using SharpCompress.Archives.Rar;
using SharpCompress.Archives.SevenZip;
using SharpCompress.Common;
using SharpCompress.Readers;
using TinfoilVibeServer.Models;
using TinfoilVibeServer.Utilities;
using ZipArchive = SharpCompress.Archives.Zip.ZipArchive;

namespace TinfoilVibeServer.Services;

public interface IArchiveHandler
{
    /// <summary>
    /// Return TitleInfo if an embedded Nintendo archive is found; otherwise null.
    /// </summary>
    IEnumerable<(string, long, NcaMetadataWithHash)> TryExtractTitleInfos(string filePath);
    
    /// <summary>
    /// Return TitleInfo if an embedded Nintendo archive is found; otherwise null.
    /// </summary>
    IEnumerable<(string, long, NcaMetadataWithHash)> TryExtractTitleInfos(Stream archiveStream, string archiveType);
}

/// <summary>
/// Tries to open a file as an archive and look for an embedded NSP/XCI.
/// The extractor is injected so that the hash of the first stream can be accessed
/// while the file is being read.
/// </summary>
public sealed class ArchiveHandler : IArchiveHandler
{
    private readonly INSPExtractor _nspExtractor;
    private readonly ILogger<ArchiveHandler> _logger;

    public ArchiveHandler(INSPExtractor nspExtractor, ILogger<ArchiveHandler> logger)
    {
        _nspExtractor = nspExtractor;
        _logger = logger;
    }

    /// <summary>
    /// Return TitleInfo if an embedded Nintendo archive is found; otherwise null.
    /// </summary>
    public IEnumerable<(string, long, NcaMetadataWithHash)> TryExtractTitleInfos(string filePath)
    {
        _logger.LogInformation("Examining archive {File} for embedded NSP", filePath);
        var ext = Path.GetExtension(filePath).ToLowerInvariant();

        try
        {
            switch (ext)
            {
                case ".zip":
                    return HandleZip(filePath);
                case ".7z":
                    return Handle7z(filePath);
                case ".rar":
                    return HandleRar(filePath);
                default:
                {
                    _logger.LogWarning("Unsupported archive type {Extension} – skipping", ext);
                    return [];
                }
            }
        }
        catch (Exception ex)
        {
            _logger.LogError("Error opening archive {File}: {Exception}", filePath, ex.Message);
            // Graceful fallback – return null
            return [];
        }
    }

    public IEnumerable<(string, long, NcaMetadataWithHash)> TryExtractTitleInfos(Stream archiveStream, string archiveType)
    {
        throw new NotImplementedException();
    }

    private IEnumerable<(string, long, NcaMetadataWithHash)> HandleZip(string path)
    {
        using var archive = ZipArchive.Open(path);
        foreach (var entry in archive.Entries)
        {
            if (entry.IsDirectory || entry.Key == null || !IsRomArchive(entry.Key)) continue;
            
            var temp = Path.GetTempFileName();
            entry.WriteToFile(temp);
            var title = _nspExtractor.ExtractFromFile(temp);   // instance call
            File.Delete(temp);
            if (title != null) yield return (entry.Key, entry.Size, title);
        }
    }

    private IEnumerable<(string, long, NcaMetadataWithHash)> Handle7z(string path)
    {
        using var archive = SevenZipArchive.Open(path);
        foreach (var entry in archive.Entries)
        {
            if (!entry.IsDirectory && entry.Key != null && IsRomArchive(entry.Key))
            {
                var temp = Path.GetTempFileName();
                entry.WriteToFile(temp);
                var title = _nspExtractor.ExtractFromFile(temp);   // instance call
                File.Delete(temp);
                if (title != null) yield return (entry.Key, entry.Size, title);
            }
        }
    }

    private IEnumerable<(string, long, NcaMetadataWithHash)> HandleRar(string path)
    {
        var titles = new List<(string, long, NcaMetadataWithHash)>();
        var entryCount = 0;
        try
        {
            using var romArchive = new RomArchiveReader(path);
            var archiveEntries = romArchive.GetEntries().ToList();
            //using var archive = RarArchive.Open(path);
            entryCount = archiveEntries.Count;
            foreach (var archiveEntry in archiveEntries)
            {
                var archiveEntryName = archiveEntry.Name;
                try
                {
                    Stream? wrapper = null;
                    
                    wrapper = new RewindableStream(archiveEntry.Stream, () =>
                    {
                        _logger.LogDebug("Rewinding archive entry {ArchiveEntry}", archiveEntryName);
                        return romArchive.GetEntries().First(romArchiveEntry => romArchiveEntry.Name == archiveEntryName).Stream;
                    },10*1024*1024, archiveEntry.Stream.Length);
                    //wrapper = new SeekableBufferedStream(archiveEntry.Stream, archiveEntry.Stream.Length, 10*1024*1024, true);
                    var title = _nspExtractor.ExtractFromStream(wrapper);
                    if (title != null)
                    {
                        titles.Add((archiveEntry.Name, archiveEntry.Stream.Length, title));
                    }
                    wrapper?.Dispose();
                }
                catch (IncompleteArchiveException incompleteArchiveException)
                {
                    _logger.LogWarning("Incomplete archive {Archive}: {Exception}", path, incompleteArchiveException.Message);
                }
                catch (Exception e)
                {
                    if (e.Message.StartsWith("Failed to extract NSP"))
                    {
                        _logger.LogError("Failed to extract title info from archive {Archive}: {Exception}", path, e.Message);
                    }
                    else if (e.Message.StartsWith("Unable to decrypt NCA section"))
                    {
                        _logger.LogError("Unable to decrypt NCA section, try updating prod.keys");
                    }
                    else
                    {
                        throw;
                    }
                }
            }
        }
        catch (Exception exception)
        {
            if (titles.Count > 0 && titles.Count == entryCount)
            {
                // broken archive but managed to read some titles
                _logger.LogInformation("Failed to fully process archive with SharpCompress, but found {Count} titles: {Exception}", titles.Count, exception.Message);
                return titles;
            }

            // Fallback to SharpSevenZip (if needed)
            _logger.LogInformation(
                "Failed to open archive with SharpCompress, falling back to SharpSevenZip {Exception}",
                exception.Message);
            using var archive = SevenZipArchive.Open(path, new ReaderOptions { LeaveStreamOpen = false });
            foreach (var entry in archive.Entries)
            {
                if (entry is not { IsDirectory: false, Key: not null } || !IsRomArchive(entry.Key)) continue;
                
                var temp = Path.GetTempFileName();
                entry.WriteToFile(temp);
                var title = _nspExtractor.ExtractFromFile(temp);   // instance call
                File.Delete(temp);
                if (title == null) continue;
                
                _logger.LogInformation("Extracted title {Key} using SharpSevenZip", entry.Key);
                titles.Add((entry.Key, entry.Size, title));
            }
        }

        return titles;
    }

    private bool IsRomArchive(string entryName)
    {
        var ext = Path.GetExtension(entryName).ToLowerInvariant();
        return ext is ".xci" or ".nsp" or ".xcz";
    }
}