using System.IO; using System.Collections.Generic; using Microsoft.Data.Sqlite; using SabreTools.Hashing; using SabreTools.Help; using SabreTools.IO; using SabreTools.IO.Extensions; namespace Headerer.Features { internal class Restore : BaseFeature { public const string Value = "Restore"; public Restore() { Name = Value; Flags.AddRange(["re", "restore"]); Description = "Restore header to file based on SHA-1"; _featureType = ParameterType.Flag; LongDescription = @"This will make use of stored copier headers and reapply them to files if they match the included hash. More than one header can be applied to a file, so they will be output to new files, suffixed with .newX, where X is a number. No input files are altered in the process. Only uncompressed files will be processed. The following systems have headers that this program can work with: - Atari 7800 - Atari Lynx - Commodore PSID Music - NEC PC - Engine / TurboGrafx 16 - Nintendo Famicom / Nintendo Entertainment System - Nintendo Famicom Disk System - Nintendo Super Famicom / Super Nintendo Entertainment System - Nintendo Super Famicom / Super Nintendo Entertainment System SPC"; // Common Features AddCommonFeatures(); AddFeature(OutputDirStringInput); } public override bool ProcessFeatures(Dictionary features) { // If the base fails, just fail out if (!base.ProcessFeatures(features)) return false; // Get only files from the inputs List files = PathTool.GetFilesOnly(Inputs); foreach (ParentablePath file in files) { RestoreHeader(file.CurrentPath, OutputDir); } return true; } /// /// Detect and replace header(s) to the given file /// /// Name of the file to be parsed /// Output directory to write the file to, empty means the same directory as the input file /// True if a header was found and appended, false otherwise public bool RestoreHeader(string file, string? outDir) { // Create the output directory if it doesn't exist if (!string.IsNullOrWhiteSpace(outDir) && !Directory.Exists(outDir)) Directory.CreateDirectory(outDir); // First, get the SHA-1 hash of the file string sha1 = HashTool.GetFileHash(file, HashType.SHA1) ?? string.Empty; // Retrieve a list of all related headers from the database List headers = RetrieveHeadersFromDatabase(sha1); // If we have nothing retrieved, we return false if (headers.Count == 0) return false; // Now loop through and create the reheadered files, if possible for (int i = 0; i < headers.Count; i++) { string outputFile = (string.IsNullOrWhiteSpace(outDir) ? $"{Path.GetFullPath(file)}.new" : Path.Combine(outDir, Path.GetFileName(file))) + i; logger.User($"Creating reheadered file: {outputFile}"); AppendBytes(file, outputFile, ByteArrayExtensions.StringToByteArray(headers[i]), null); logger.User("Reheadered file created!"); } return true; } /// /// Retrieve headers from the database /// /// SHA-1 of the deheadered file /// List of strings representing the headers to add private List RetrieveHeadersFromDatabase(string SHA1) { // Ensure the database exists EnsureDatabase(); // Open the database connection SqliteConnection dbc = new SqliteConnection(HeadererConnectionString); dbc.Open(); // Create the output list of headers List headers = []; string query = $"SELECT header, type FROM data WHERE sha1='{SHA1}'"; SqliteCommand slc = new SqliteCommand(query, dbc); SqliteDataReader sldr = slc.ExecuteReader(); if (sldr.HasRows) { while (sldr.Read()) { logger.Verbose($"Found match with rom type '{sldr.GetString(1)}'"); headers.Add(sldr.GetString(0)); } } else { logger.Warning("No matching header could be found!"); } // Dispose of database objects slc.Dispose(); sldr.Dispose(); dbc.Dispose(); return headers; } /// /// Add an aribtrary number of bytes to the inputted file /// /// File to be appended to /// Outputted file /// Bytes to be added to head of file /// Bytes to be added to tail of file private static void AppendBytes(string input, string output, byte[]? bytesToAddToHead, byte[]? bytesToAddToTail) { // If any of the inputs are invalid, skip if (!File.Exists(input)) return; using FileStream fsr = File.OpenRead(input); using FileStream fsw = File.OpenWrite(output); AppendBytes(fsr, fsw, bytesToAddToHead, bytesToAddToTail); } /// /// Add an aribtrary number of bytes to the inputted stream /// /// Stream to be appended to /// Outputted stream /// Bytes to be added to head of stream /// Bytes to be added to tail of stream private static void AppendBytes(Stream input, Stream output, byte[]? bytesToAddToHead, byte[]? bytesToAddToTail) { // Write out prepended bytes if (bytesToAddToHead != null && bytesToAddToHead.Length > 0) output.Write(bytesToAddToHead, 0, bytesToAddToHead.Length); // Now copy the existing file over input.CopyTo(output); // Write out appended bytes if (bytesToAddToTail != null && bytesToAddToTail.Length > 0) output.Write(bytesToAddToTail, 0, bytesToAddToTail.Length); } } }