using System;
using System.Collections.Generic;
using System.IO;
using Microsoft.Data.Sqlite;
using SabreTools.Hashing;
using SabreTools.IO.Extensions;
namespace Headerer
{
internal static class Restore
{
///
/// 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 static 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;
Console.WriteLine($"Creating reheadered file: {outputFile}");
AppendBytes(file, outputFile, ByteArrayExtensions.StringToByteArray(headers[i]), null);
Console.WriteLine("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 static List RetrieveHeadersFromDatabase(string SHA1)
{
// Ensure the database exists
Database.EnsureDatabase();
// Open the database connection
var dbc = new SqliteConnection(Database.HeadererConnectionString);
dbc.Open();
// Create the output list of headers
List headers = [];
string query = $"SELECT header, type FROM data WHERE sha1='{SHA1}'";
var slc = new SqliteCommand(query, dbc);
SqliteDataReader sldr = slc.ExecuteReader();
if (sldr.HasRows)
{
while (sldr.Read())
{
Console.WriteLine($"Found match with rom type '{sldr.GetString(1)}'"); // TODO: Gate behind debug flag
headers.Add(sldr.GetString(0));
}
}
else
{
Console.Error.WriteLine("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);
}
}
}