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

using Microsoft.Data.Sqlite;
using ZB.MOM.WW.Secrets.Abstractions;
using ZB.MOM.WW.Secrets.Cli.Interactive;
using ZB.MOM.WW.Secrets.Crypto;
using ZB.MOM.WW.Secrets.MasterKey;
using ZB.MOM.WW.Secrets.Rotation;
using ZB.MOM.WW.Secrets.Sqlite;

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

/// <summary>
/// Exercises <see cref="KekDoctor"/> against the REAL SQLite store and REAL AES-256-GCM cipher — the
/// lockout-triage path an operator runs to learn, per row, which KEK it is wrapped under and whether
/// the session KEK actually opens it, then remedies via a guided rewrap-all. KEK-A and KEK-B are two
/// distinct 32-byte <see cref="LiteralMasterKeyProvider"/> keys (distinct derived <c>kek_id</c>s).
/// </summary>
public sealed class KekDoctorTests : IAsyncLifetime, IDisposable
{
    private readonly string _dbPath =
        Path.Combine(Path.GetTempPath(), $"zb-secrets-doctor-{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).
    private readonly LiteralMasterKeyProvider _kekA =
        new(Convert.ToBase64String(FillKey(0xA1)));
    private readonly LiteralMasterKeyProvider _kekB =
        new(Convert.ToBase64String(FillKey(0xB2)));

    public KekDoctorTests()
    {
        _factory = new SecretsSqliteConnectionFactory(_dbPath);
        _store = new SqliteSecretStore(_factory);
    }

    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;
    }

    // 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",
    };

    // A degraded session: no KEK, no cipher (metadata-only).
    private SecretsSession DegradedSession() => new()
    {
        Target = new TargetConfigReader().Manual(_dbPath, new MasterKeyOptions()),
        Store = _store,
        Migrator = new SqliteSecretsStoreMigrator(_factory),
        MasterKey = null,
        Cipher = null,
        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);
    }

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

        KekDoctorReport report = await new KekDoctor().DiagnoseAsync(Session(_kekA), CancellationToken.None);

        Assert.True(report.Healthy);
        Assert.Equal(_kekA.KekId, report.SessionKekId);
        Assert.Equal(2, report.Rows.Count);
        Assert.Equal(2, report.Total); // no rows vanished → Total equals Rows.Count.
        Assert.All(report.Rows, r => Assert.Equal(RowKekStatus.Ok, r.Status));
        // Rows ordered by name: "ldap/b" precedes "sql/a".
        Assert.Equal(["ldap/b", "sql/a"], report.Rows.Select(r => r.SecretName));
    }

    [Fact]
    public async Task Diagnose_verdict_reads_fresh_row_kek_not_stale_metadata()
    {
        // Seed under A; a first diagnosis on session A reports Ok.
        await SeedAsync("sql/a", "value-a", _kekA);
        KekDoctorReport first = await new KekDoctor().DiagnoseAsync(Session(_kekA), CancellationToken.None);
        Assert.Equal(RowKekStatus.Ok, Assert.Single(first.Rows).Status);

        // Re-seal the SAME row under B (simulates a concurrent re-wrap/re-set moving its KEK). A second
        // diagnosis on session A must classify from the FRESH row.KekId (B) → WrongKek surfacing B's id,
        // never a false Corrupt from probing a row the session cannot open.
        await SeedAsync("sql/a", "value-a", _kekB);
        KekDoctorReport second = await new KekDoctor().DiagnoseAsync(Session(_kekA), CancellationToken.None);

        KekDiagnosis row = Assert.Single(second.Rows);
        Assert.Equal(RowKekStatus.WrongKek, row.Status);
        Assert.Equal(_kekB.KekId, row.RowKekId);
    }

    [Fact]
    public async Task Diagnose_reports_wrong_kek_rows_with_their_kekid()
    {
        // Rows sealed under A; the session runs on B → the operator must learn A is the KEK to hunt.
        await SeedAsync("sql/a", "value-a", _kekA);

        KekDoctorReport report = await new KekDoctor().DiagnoseAsync(Session(_kekB), CancellationToken.None);

        Assert.False(report.Healthy);
        Assert.Equal(_kekB.KekId, report.SessionKekId);
        KekDiagnosis row = Assert.Single(report.Rows);
        Assert.Equal(RowKekStatus.WrongKek, row.Status);
        Assert.Equal(_kekA.KekId, row.RowKekId);
    }

    [Fact]
    public async Task Diagnose_reports_corrupt_row_when_kekid_matches_but_unwrap_fails()
    {
        await SeedAsync("sql/a", "value-a", _kekA);

        // Tamper one byte of the DEK wrap tag: kek_id still matches the session KEK, but the unwrap
        // (and therefore the decrypt probe) now fails closed. UpsertAsync bumps revision — irrelevant.
        StoredSecret row = (await _store.GetAsync(new SecretName("sql/a"), CancellationToken.None))!;
        byte[] tag = (byte[])row.WrapTag.Clone();
        tag[0] ^= 0xFF;
        await _store.UpsertAsync(row with { WrapTag = tag }, CancellationToken.None);

        KekDoctorReport report = await new KekDoctor().DiagnoseAsync(Session(_kekA), CancellationToken.None);

        Assert.False(report.Healthy);
        KekDiagnosis diag = Assert.Single(report.Rows);
        Assert.Equal(RowKekStatus.Corrupt, diag.Status);
        Assert.Equal(_kekA.KekId, diag.RowKekId);
    }

    [Fact]
    public async Task Diagnose_includes_tombstoned_rows()
    {
        await SeedAsync("sql/live", "live", _kekA);
        await SeedAsync("sql/dead", "dead", _kekA);
        await _store.DeleteAsync(new SecretName("sql/dead"), "carol", CancellationToken.None);

        KekDoctorReport report = await new KekDoctor().DiagnoseAsync(Session(_kekA), CancellationToken.None);

        // A tombstoned row still appears — it blocks rewrap-all (it carries a KEK-wrapped DEK).
        Assert.Equal(2, report.Rows.Count);
        Assert.Contains(report.Rows, r => r.SecretName == "sql/dead");
    }

    [Fact]
    public async Task Diagnose_throws_on_degraded_session()
    {
        InvalidOperationException ex = await Assert.ThrowsAsync<InvalidOperationException>(
            () => new KekDoctor().DiagnoseAsync(DegradedSession(), CancellationToken.None));

        Assert.Contains("key", ex.Message, StringComparison.OrdinalIgnoreCase);
    }

    [Fact]
    public async Task RewrapAll_throws_on_degraded_session()
    {
        InvalidOperationException ex = await Assert.ThrowsAsync<InvalidOperationException>(
            () => new KekDoctor().RewrapAllAsync(DegradedSession(), _kekA, CancellationToken.None));

        Assert.Contains("key", ex.Message, StringComparison.OrdinalIgnoreCase);
    }

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

        // Session runs on B; rows are on A → rewrap A onto the session KEK (B).
        SecretsSession session = Session(_kekB);
        RewrapReport report = await new KekDoctor().RewrapAllAsync(session, _kekA, CancellationToken.None);

        Assert.Equal(2, report.Total);
        Assert.Equal(2, report.Rewrapped);
        Assert.Equal(0, report.AlreadyCurrent);

        // Rows now decrypt under the session KEK (B), and the doctor re-diagnoses healthy.
        var cipherB = new AesGcmEnvelopeCipher(_kekB);
        StoredSecret a = (await _store.GetAsync(new SecretName("sql/a"), CancellationToken.None))!;
        Assert.Equal(_kekB.KekId, a.KekId);
        Assert.Equal("value-a", cipherB.Decrypt(a));

        KekDoctorReport after = await new KekDoctor().DiagnoseAsync(session, CancellationToken.None);
        Assert.True(after.Healthy);
    }

    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.
            }
        }
    }
}