scadaproj/ZB.MOM.WW.Secrets/tests/ZB.MOM.WW.Secrets.Tests/Rotation/KekRotationServiceTests.cs

using Microsoft.Data.Sqlite;
using ZB.MOM.WW.Secrets.Abstractions;
using ZB.MOM.WW.Secrets.Crypto;
using ZB.MOM.WW.Secrets.Rotation;
using ZB.MOM.WW.Secrets.Sqlite;
using ZB.MOM.WW.Secrets.Tests.Fakes;

namespace ZB.MOM.WW.Secrets.Tests.Rotation;

/// <summary>
/// End-to-end KEK-rotation tests against the REAL SQLite store and REAL AES-256-GCM cipher — the
/// combination that actually proves a re-wrapped row decrypts under the new KEK (and no longer
/// under the old one).
/// </summary>
public sealed class KekRotationServiceTests : IAsyncLifetime, IDisposable
{
    private readonly string _dbPath =
        Path.Combine(Path.GetTempPath(), $"zb-secrets-rotate-{Guid.NewGuid():N}.db");

    private readonly SecretsSqliteConnectionFactory _factory;
    private readonly SqliteSecretStore _store;

    private readonly FakeMasterKeyProvider _oldKek = new("kek-old");
    private readonly FakeMasterKeyProvider _newKek = new("kek-new");

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

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

    public Task DisposeAsync() => Task.CompletedTask;

    // Seeds a secret sealed under a given KEK provider (defaults to the old KEK).
    private async Task SeedAsync(string name, string value, IMasterKeyProvider? kek = null)
    {
        var cipher = new AesGcmEnvelopeCipher(kek ?? _oldKek);
        StoredSecret row = cipher.Encrypt(new SecretName(name), value, SecretContentType.Text);
        await _store.UpsertAsync(row, CancellationToken.None);
    }

    private KekRotationService Service() => new(_store, new AesGcmEnvelopeCipher(_oldKek));

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

        RewrapReport report = await Service().RewrapAllAsync(_oldKek, _newKek, CancellationToken.None);

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

        var newCipher = new AesGcmEnvelopeCipher(_newKek);
        var oldCipher = new AesGcmEnvelopeCipher(_oldKek);
        foreach ((string name, string value) in new[] { ("sql/a", "value-a"), ("sql/b", "value-b"), ("ldap/c", "value-c") })
        {
            StoredSecret row = (await _store.GetAsync(new SecretName(name), CancellationToken.None))!;
            Assert.Equal("kek-new", row.KekId);
            // Decrypts cleanly under the NEW KEK...
            Assert.Equal(value, newCipher.Decrypt(row));
            // ...and fails closed under the OLD KEK (kek_id no longer matches).
            Assert.Throws<SecretDecryptionException>(() => oldCipher.Decrypt(row));
        }
    }

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

        RewrapReport first = await Service().RewrapAllAsync(_oldKek, _newKek, CancellationToken.None);
        Assert.Equal(2, first.Rewrapped);

        RewrapReport second = await Service().RewrapAllAsync(_oldKek, _newKek, CancellationToken.None);
        Assert.Equal(2, second.Total);
        Assert.Equal(0, second.Rewrapped);
        Assert.Equal(2, second.AlreadyCurrent);
    }

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

        RewrapReport report = await Service().RewrapAllAsync(_oldKek, _newKek, CancellationToken.None);

        Assert.Equal(2, report.Rewrapped);
        StoredSecret dead = (await _store.GetAsync(new SecretName("sql/dead"), CancellationToken.None))!;
        Assert.Equal("kek-new", dead.KekId);
        Assert.True(dead.IsDeleted);
    }

    [Fact]
    public async Task RewrapAll_DoesNotBumpRevisionOrUpdatedUtc()
    {
        await SeedAsync("sql/a", "value-a");
        StoredSecret before = (await _store.GetAsync(new SecretName("sql/a"), CancellationToken.None))!;

        await Service().RewrapAllAsync(_oldKek, _newKek, CancellationToken.None);

        StoredSecret after = (await _store.GetAsync(new SecretName("sql/a"), CancellationToken.None))!;
        Assert.Equal(before.Revision, after.Revision);
        Assert.Equal(before.UpdatedUtc, after.UpdatedUtc);
    }

    [Fact]
    public async Task RewrapAll_IdenticalKekIds_Throws()
    {
        await SeedAsync("sql/a", "value-a");
        var same = new FakeMasterKeyProvider("kek-old");

        await Assert.ThrowsAsync<ArgumentException>(
            () => Service().RewrapAllAsync(_oldKek, same, CancellationToken.None));
    }

    [Fact]
    public async Task RewrapAll_RowOnUnknownKek_Aborts_ButKeepsPriorProgress()
    {
        // Names are scanned in ORDER BY name: "a-old" and "b-old" migrate, then "z-third" (a KEK that
        // is neither old nor new) aborts the pass.
        await SeedAsync("a-old", "va");
        await SeedAsync("b-old", "vb");
        await SeedAsync("z-third", "vz", new FakeMasterKeyProvider("kek-third"));

        await Assert.ThrowsAsync<SecretDecryptionException>(
            () => Service().RewrapAllAsync(_oldKek, _newKek, CancellationToken.None));

        // The two old-KEK rows scanned before the anomaly are already persisted on the new KEK,
        // so a re-run (after removing the anomaly) resumes cleanly.
        Assert.Equal("kek-new", (await _store.GetAsync(new SecretName("a-old"), CancellationToken.None))!.KekId);
        Assert.Equal("kek-new", (await _store.GetAsync(new SecretName("b-old"), CancellationToken.None))!.KekId);
        Assert.Equal("kek-third", (await _store.GetAsync(new SecretName("z-third"), CancellationToken.None))!.KekId);
    }

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