using System.IO.Compression; using System.Security.Cryptography; using System.Text.Json; using AIStudio.Tools.PluginSystem; using SettingsVersion = AIStudio.Settings.Version; namespace AIStudio.Tools.DataBackup; /// /// Prepares a restore: unpacks a backup next to the data and configuration folders, so that the next /// start can put it in place, cf. DataRestoreFormat. /// /// /// Nothing in use changes here. While AI Studio runs, its files are open, and on Windows, they cannot be /// replaced at all. The runtime replaces them at the next start, before the app opens any of them, and /// rolls back when that fails. /// public static class DataRestorePreparer { private const int BUFFER_SIZE = 64 * 1024; private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataRestorePreparer).Namespace, nameof(DataRestorePreparer)); /// /// A file of the backup to unpack, and where it goes. /// private sealed record PendingFile(ZipArchiveEntry Entry, DataBackupManifestEntry Listed, string DestinationPath); /// /// Unpacks a backup for the next start. /// /// /// Every file has to match the size and the hash the manifest lists for it: a restore replaces all /// data, so a damaged backup must not get that far. A file the manifest does not list, or one /// which belongs to no element the restore replaces, stays in the backup. A restore prepared /// before is discarded first; one which fails or is cancelled leaves nothing behind. /// /// The backup. /// Where the restore goes. /// Receives how many files are unpacked. /// Cancels the preparation. /// Whether the restore is prepared, or why not. public static async Task PrepareAsync(string archivePath, DataRestoreTarget target, IProgress? progress, CancellationToken token) { Discard(target); try { await using var archive = await ZipFile.OpenReadAsync(archivePath, token); var manifestResult = await DataBackupReader.ReadManifestAsync(archive, token); if (manifestResult.Manifest is not { } manifest) return new(false, false, manifestResult.Issue); var versionIssue = CheckVersions(manifest, target); if (!string.IsNullOrWhiteSpace(versionIssue)) return new(false, false, versionIssue); var elements = DetermineElements(manifest, target); var files = DetermineFiles(archive, manifest, elements, target); progress?.Report(new(0, files.Count)); var buffer = new byte[BUFFER_SIZE]; for (var index = 0; index < files.Count; index++) { token.ThrowIfCancellationRequested(); if (!await UnpackAsync(files[index], buffer, token)) { Discard(target); return new(false, false, string.Format(TB("The backup is damaged: the file '{0}' does not match its content list."), files[index].Listed.Name)); } progress?.Report(new(index + 1, files.Count)); } var pendingDirectory = PendingDirectory(target.DataDirectory); Directory.CreateDirectory(pendingDirectory); var plan = new DataRestorePlan(DataRestoreFormat.PLAN_FORMAT_VERSION, DateTimeOffset.UtcNow, manifest.CreatedAtUtc, manifest.AppVersion, elements); await File.WriteAllBytesAsync(Path.Join(pendingDirectory, DataRestoreFormat.PLAN_FILE_NAME), JsonSerializer.SerializeToUtf8Bytes(plan, DataBackupFormat.JSON_OPTIONS), token); // Last of all: from now on, the next start applies the restore. await File.WriteAllTextAsync(Path.Join(pendingDirectory, DataRestoreFormat.READY_FILE_NAME), plan.PreparedAtUtc.ToString("O"), token); return new(true, false, string.Empty); } catch (InvalidDataException) { Discard(target); return new(false, false, TB("This file is no backup of AI Studio, or it was not written completely.")); } catch { Discard(target); throw; } } /// /// Removes a prepared restore, complete or not, so that the next start leaves everything as it is. /// /// /// A restore which the runtime has not finished stays: its journal is the only way back to the data /// it replaced, cf. DataRestoreFollowUp.IsRestoreUnfinished. /// /// Where the restore was to go. public static void Discard(DataRestoreTarget target) { if (DataRestoreFollowUp.IsRestoreUnfinished(target.DataDirectory)) throw new InvalidOperationException("The runtime has not finished the last restore yet, so it cannot be discarded."); // First of all: should removing the rest fail halfway, the next start must not apply what is left. var readyPath = Path.Join(PendingDirectory(target.DataDirectory), DataRestoreFormat.READY_FILE_NAME); if (File.Exists(readyPath)) File.Delete(readyPath); foreach (var pendingDirectory in new[] { PendingDirectory(target.DataDirectory), PendingDirectory(target.ConfigDirectory) }) if (Directory.Exists(pendingDirectory)) Directory.Delete(pendingDirectory, true); } /// /// Whether a restore is prepared completely and waits for the next start. /// /// Where the restore goes. /// True when the next start applies a restore. public static bool IsPrepared(DataRestoreTarget target) => DataRestoreFollowUp.IsRestorePending(target.DataDirectory); private static string PendingDirectory(string directory) => DataRestoreFormat.DirectoryIn(directory, DataRestoreFormat.PENDING_DIRECTORY); /// /// Refuses a backup of a later version of AI Studio. Its settings or its index may have a form /// this version cannot read, and it would refuse to write them, or break them. /// /// Why the backup cannot be restored, or an empty string. private static string CheckVersions(DataBackupManifest manifest, DataRestoreTarget target) { var settingsAreNewer = !Enum.TryParse(manifest.SettingsVersion, out var settingsVersion) || settingsVersion > target.SettingsVersion; var appIsNewer = Version.TryParse(manifest.AppVersion, out var backupAppVersion) && Version.TryParse(target.AppVersion, out var appVersion) && backupAppVersion > appVersion; return settingsAreNewer || appIsNewer ? TB("This backup was made by a newer version of AI Studio. Please update AI Studio first.") : string.Empty; } /// /// Determines the elements the restore replaces: those of the list of the backup contents which the /// backup holds anything for. /// /// /// An element the backup holds nothing for stays as it is, since nothing tells whether it was empty /// or whether the version of AI Studio which made the backup knew it at all. The indexes and the /// media transcripts are replaced only when the backup holds the indexes, cf. /// DataBackupManifest.IncludesDataSourceIndexes and DataBackupEntry.ComesFromIndexCopies. /// private static List DetermineElements(DataBackupManifest manifest, DataRestoreTarget target) { var elements = new List(); foreach (var entry in DataBackupContents.ENTRIES.Where(entry => entry.IsIncluded)) { if (entry.ComesFromIndexCopies && !manifest.IncludesDataSourceIndexes) continue; var entryName = DataBackupFormat.EntryName(entry.Location, entry.RelativePath); if (!manifest.Entries.Any(listed => DataBackupContents.Covers(entryName, listed.Name))) continue; var keep = DataBackupContents.ENTRIES .Where(excluded => !excluded.IsIncluded && excluded.Location == entry.Location && entry.IsFolder && DataBackupContents.Covers(entry.RelativePath, excluded.RelativePath)) .Select(excluded => excluded.RelativePath) .Concat(entry.Kind is DataBackupEntryKind.PLUGINS ? target.OrganizationPluginFolders : []) .ToList(); elements.Add(new(entry.Location, entry.RelativePath, keep)); } return elements; } /// /// Determines the files to unpack: those the manifest lists, which belong to an element the restore /// replaces, and which the element does not keep. /// private static List DetermineFiles(ZipArchive archive, DataBackupManifest manifest, IReadOnlyList elements, DataRestoreTarget target) { var listedEntries = manifest.Entries.ToDictionary(listed => listed.Name, StringComparer.Ordinal); var files = new List(); foreach (var entry in archive.Entries) { var entryName = ZipEntryNames.Normalize(entry.FullName); if (!listedEntries.TryGetValue(entryName, out var listed)) continue; foreach (var element in elements) { var root = DataBackupFormat.EntryName(element.Location, string.Empty); if (!entryName.StartsWith(root, StringComparison.Ordinal)) continue; var relativePath = entryName[root.Length..]; if (!DataBackupContents.Covers(element.RelativePath, relativePath) || element.Keep.Any(kept => DataBackupContents.Covers(kept, relativePath))) continue; var directory = element.Location is DataBackupLocation.CONFIG_DIRECTORY ? target.ConfigDirectory : target.DataDirectory; files.Add(new(entry, listed, ZipEntryNames.GetDestinationPath(PendingDirectory(directory), relativePath))); break; } } return files; } /// /// Unpacks one file, and checks its size and hash while it streams through. /// /// False when the file does not match the manifest. private static async Task UnpackAsync(PendingFile file, byte[] buffer, CancellationToken token) { Directory.CreateDirectory(Path.GetDirectoryName(file.DestinationPath)!); using var hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256); long size = 0; await using (var source = await file.Entry.OpenAsync(token)) await using (var destination = new FileStream(file.DestinationPath, FileMode.CreateNew, FileAccess.Write, FileShare.None, BUFFER_SIZE, FileOptions.Asynchronous | FileOptions.SequentialScan)) { int read; while ((read = await source.ReadAsync(buffer, token)) > 0) { hash.AppendData(buffer, 0, read); await destination.WriteAsync(buffer.AsMemory(0, read), token); size += read; } } if (size != file.Listed.Size || Convert.ToHexStringLower(hash.GetHashAndReset()) != file.Listed.Sha256) return false; // The time of the last change tells, e.g., the order of the chats in a workspace: File.SetLastWriteTime(file.DestinationPath, file.Entry.LastWriteTime.LocalDateTime); return true; } }