scadaproj/ZB.MOM.WW.Secrets/tests/ZB.MOM.WW.Secrets.Tests/Cli/Interactive/Flows/KekDoctorFlowTests.cs

using Microsoft.Data.Sqlite;
using Spectre.Console.Testing;
using ZB.MOM.WW.Secrets.Abstractions;
using ZB.MOM.WW.Secrets.Cli.Interactive;
using ZB.MOM.WW.Secrets.Cli.Interactive.Flows;
using ZB.MOM.WW.Secrets.Crypto;
using ZB.MOM.WW.Secrets.MasterKey;
using ZB.MOM.WW.Secrets.Sqlite;

namespace ZB.MOM.WW.Secrets.Tests.Cli.Interactive.Flows;

/// <summary>
/// Drives the lockout-recovery <see cref="KekDoctorFlow"/> against a REAL temp-SQLite store and REAL
/// AES-256-GCM cipher. KEK-A and KEK-B are two distinct 32-byte <see cref="LiteralMasterKeyProvider"/>
/// keys (distinct derived <c>kek_id</c>s): rows are sealed under one and the session opened under the
/// other to reproduce the wrong-KEK lockout. Each flow renders to a scripted <see cref="TestConsole"/>
/// so the tests assert both the rendered <see cref="TestConsole.Output"/> and the resulting store state.
/// The invariant is that the operator recovers a locked-out store from one screen — either by re-wrapping
/// from the recovered old KEK or by re-setting the affected secrets — without any key or value ever
/// reaching the terminal.
/// </summary>
public sealed class KekDoctorFlowTests : IAsyncLifetime, IDisposable
{
    private readonly string _dbPath =
        Path.Combine(Path.GetTempPath(), $"zb-secrets-doctorflow-{Guid.NewGuid():N}.db");

    private readonly SecretsSqliteConnectionFactory _factory;
    private readonly SqliteSecretStore _store;

    // Two distinct 32-byte keys → two distinct derived kek_ids (KEK-A / KEK-B). The base64 forms are
    // retained so a test can paste KEK-A at the prompt and assert it never echoes back to the terminal.
    private readonly string _kekABase64 = Convert.ToBase64String(FillKey(0xA1));
    private readonly string _kekBBase64 = Convert.ToBase64String(FillKey(0xB2));
    private readonly LiteralMasterKeyProvider _kekA;
    private readonly LiteralMasterKeyProvider _kekB;

    public KekDoctorFlowTests()
    {
        _factory = new SecretsSqliteConnectionFactory(_dbPath);
        _store = new SqliteSecretStore(_factory);
        _kekA = new LiteralMasterKeyProvider(_kekABase64);
        _kekB = new LiteralMasterKeyProvider(_kekBBase64);
    }

    public async Task InitializeAsync() =>
        await new SqliteSecretsStoreMigrator(_factory).MigrateAsync(CancellationToken.None);

    public Task DisposeAsync() => Task.CompletedTask;

    private static byte[] FillKey(byte b)
    {
        byte[] key = new byte[32];
        Array.Fill(key, b);
        return key;
    }

    private static TestConsole NewConsole()
    {
        var console = (TestConsole)new TestConsole().Interactive();
        console.Profile.Width = 240; // wide enough that the summary lines never truncate a kek_id
        return console;
    }

    // A live session over the shared store, keyed by the supplied provider.
    private SecretsSession Session(IMasterKeyProvider kek) => new()
    {
        Target = new TargetConfigReader().Manual(_dbPath, new MasterKeyOptions()),
        Store = _store,
        Migrator = new SqliteSecretsStoreMigrator(_factory),
        MasterKey = kek,
        Cipher = new AesGcmEnvelopeCipher(kek),
        StoreKind = "sqlite",
    };

    private async Task SeedAsync(string name, string value, IMasterKeyProvider kek)
    {
        var cipher = new AesGcmEnvelopeCipher(kek);
        StoredSecret row = cipher.Encrypt(new SecretName(name), value, SecretContentType.Text);
        await _store.UpsertAsync(row, CancellationToken.None);
    }

    private async Task<StoredSecret> GetAsync(string name) =>
        (await _store.GetAsync(new SecretName(name), CancellationToken.None))!;

    [Fact]
    public async Task Healthy_store_renders_all_ok_summary()
    {
        await SeedAsync("svc/a", "value-a", _kekA);
        await SeedAsync("svc/b", "value-b", _kekA);

        TestConsole console = NewConsole();
        await new KekDoctorFlow().RunAsync(console, Session(_kekA), CancellationToken.None);

        string output = console.Output;
        Assert.Contains(_kekA.KekId, output, StringComparison.Ordinal); // the session KEK is surfaced
        Assert.Contains("All rows open", output, StringComparison.Ordinal); // green all-clear
        Assert.DoesNotContain("Rewrap from old KEK", output, StringComparison.Ordinal); // no remedy menu

        Assert.True(new KekDoctorFlow().RequiresKek);
        Assert.Equal("KEK doctor (lockout recovery)", new KekDoctorFlow().Title);
    }

    [Fact]
    public async Task Wrong_kek_rows_offer_rewrap_and_rewrap_succeeds()
    {
        // Rows sealed under KEK-A; the session runs on KEK-B → both rows are wrong-KEK.
        await SeedAsync("svc/a", "value-a", _kekA);
        await SeedAsync("svc/b", "value-b", _kekA);
        SecretsSession session = Session(_kekB);

        TestConsole console = NewConsole();
        console.Input.PushKey(ConsoleKey.Enter); // remedy menu: "Rewrap from old KEK" (index 0)
        console.Input.PushKey(ConsoleKey.DownArrow); // old-key source: skip "Environment variable"
        console.Input.PushKey(ConsoleKey.DownArrow); // skip "Key file"
        console.Input.PushKey(ConsoleKey.Enter); // choose "Paste base64 key" (index 2)
        console.Input.PushTextWithEnter(_kekABase64); // paste KEK-A material (masked)
        console.Input.PushTextWithEnter("y"); // confirm the rewrap (default is false)

        await new KekDoctorFlow().RunAsync(console, session, CancellationToken.None);

        // Both rows now decrypt under the session KEK (B).
        var cipherB = new AesGcmEnvelopeCipher(_kekB);
        StoredSecret a = await GetAsync("svc/a");
        StoredSecret b = await GetAsync("svc/b");
        Assert.Equal(_kekB.KekId, a.KekId);
        Assert.Equal(_kekB.KekId, b.KekId);
        Assert.Equal("value-a", cipherB.Decrypt(a));
        Assert.Equal("value-b", cipherB.Decrypt(b));

        Assert.Contains("Re-wrapped", console.Output, StringComparison.Ordinal);
        Assert.Contains("2", console.Output, StringComparison.Ordinal); // the rewrap count
        Assert.DoesNotContain(_kekABase64, console.Output, StringComparison.Ordinal); // pasted key never echoed
    }

    [Fact]
    public async Task Rewrap_requires_explicit_confirmation()
    {
        await SeedAsync("svc/a", "value-a", _kekA);
        await SeedAsync("svc/b", "value-b", _kekA);
        SecretsSession session = Session(_kekB);

        TestConsole console = NewConsole();
        console.Input.PushKey(ConsoleKey.Enter); // "Rewrap from old KEK"
        console.Input.PushKey(ConsoleKey.DownArrow);
        console.Input.PushKey(ConsoleKey.DownArrow);
        console.Input.PushKey(ConsoleKey.Enter); // "Paste base64 key"
        console.Input.PushTextWithEnter(_kekABase64);
        console.Input.PushTextWithEnter("n"); // DECLINE the confirmation

        await new KekDoctorFlow().RunAsync(console, session, CancellationToken.None);

        // Nothing was re-wrapped: both rows remain under KEK-A.
        Assert.Equal(_kekA.KekId, (await GetAsync("svc/a")).KekId);
        Assert.Equal(_kekA.KekId, (await GetAsync("svc/b")).KekId);
    }

    [Fact]
    public async Task Lost_kek_path_offers_reset_of_affected_secrets()
    {
        // Two wrong-KEK rows (sealed under KEK-A; session on KEK-B). Ordinal order: svc/a precedes svc/b.
        await SeedAsync("svc/a", "value-a", _kekA);
        await SeedAsync("svc/b", "value-b", _kekA);
        SecretsSession session = Session(_kekB);

        TestConsole console = NewConsole();
        console.Input.PushKey(ConsoleKey.DownArrow); // remedy menu: move to
        console.Input.PushKey(ConsoleKey.Enter);     // "Old KEK is lost — re-set affected secrets" (index 1)
        console.Input.PushTextWithEnter("y");        // svc/a: accept re-set
        console.Input.PushTextWithEnter("NEW-A-VALUE"); // masked new value
        console.Input.PushKey(ConsoleKey.Enter);     // content type: Text (first choice)
        console.Input.PushTextWithEnter("n");        // svc/b: decline re-set

        await new KekDoctorFlow().RunAsync(console, session, CancellationToken.None);

        // svc/a was re-set under the session KEK (B) with the new value.
        var cipherB = new AesGcmEnvelopeCipher(_kekB);
        StoredSecret a = await GetAsync("svc/a");
        Assert.Equal(_kekB.KekId, a.KekId);
        Assert.Equal("NEW-A-VALUE", cipherB.Decrypt(a));

        // svc/b was left untouched → still wrong-KEK (under KEK-A).
        Assert.Equal(_kekA.KekId, (await GetAsync("svc/b")).KekId);

        Assert.DoesNotContain("NEW-A-VALUE", console.Output, StringComparison.Ordinal); // value never echoed
    }

    public void Dispose()
    {
        SqliteConnection.ClearAllPools();
        foreach (string path in new[] { _dbPath, _dbPath + "-wal", _dbPath + "-shm" })
        {
            try
            {
                if (File.Exists(path))
                {
                    File.Delete(path);
                }
            }
            catch (IOException)
            {
                // Best-effort temp cleanup.
            }
        }
    }
}