using NSec.Cryptography;

namespace DodoSSH.Crypto.Tests;

/// <summary>
/// Vault key grants, per-item data keys, and the AAD binding that holds them in place.
/// </summary>
/// <remarks>
/// Most of these are negative. The AAD binding is the most valuable structural property in the
/// design — it is what stops a server that holds every ciphertext from moving one between rows,
/// rolling a row back, or replaying a revoked grant — and a property like that is only demonstrated
/// by the substitutions that fail.
/// </remarks>
public sealed class VaultAndItemKeyTests
{
    private static readonly Guid VaultA = Guid.Parse("0192f0c8-1111-7c3d-8e4f-5a6b7c8d9e0f");
    private static readonly Guid VaultB = Guid.Parse("0192f0c8-2222-7c3d-8e4f-5a6b7c8d9e0f");
    private static readonly Guid ItemA = Guid.Parse("0192f0c8-3333-7c3d-8e4f-5a6b7c8d9e0f");
    private static readonly Guid ItemB = Guid.Parse("0192f0c8-4444-7c3d-8e4f-5a6b7c8d9e0f");
    private static readonly Guid DataKeyId = Guid.Parse("0192f0c8-5555-7c3d-8e4f-5a6b7c8d9e0f");

    // ---- Vault keys ----

    [Fact]
    public void AVaultKey_IsThirtyTwoRandomBytes()
    {
        var first = VaultKeys.Create();
        var second = VaultKeys.Create();

        first.Length.ShouldBe(CryptoSpec.SymmetricKeySize);
        first.ShouldNotBe(second);
    }

    [Fact]
    public void AGrant_OpensForItsRecipient()
    {
        using var member = Key.Create(KeyAgreementAlgorithm.X25519);
        var vaultKey = VaultKeys.Create();

        var grant = VaultKeys.WrapTo(
            vaultKey, member.PublicKey.Export(KeyBlobFormat.RawPublicKey), VaultA, 1);

        VaultKeys.TryUnwrap(member, grant, VaultA, 1).ShouldBe(vaultKey);
    }

    [Fact]
    public void AGrantForAnotherVault_DoesNotOpen()
    {
        using var member = Key.Create(KeyAgreementAlgorithm.X25519);
        var vaultKey = VaultKeys.Create();

        var grant = VaultKeys.WrapTo(
            vaultKey, member.PublicKey.Export(KeyBlobFormat.RawPublicKey), VaultA, 1);

        VaultKeys.TryUnwrap(member, grant, VaultB, 1).ShouldBeNull();
    }

    [Fact]
    public void AGrantFromASupersededGeneration_DoesNotOpen()
    {
        // This is what makes a rekey stick against a malicious server: it cannot re-serve the old
        // generation's grant to a removed member. Bounded to future reads only — whatever they
        // already downloaded is already gone, which is why offboarding means rotating the SSH
        // credentials themselves.
        using var member = Key.Create(KeyAgreementAlgorithm.X25519);
        var vaultKey = VaultKeys.Create();

        var grant = VaultKeys.WrapTo(
            vaultKey, member.PublicKey.Export(KeyBlobFormat.RawPublicKey), VaultA, 1);

        VaultKeys.TryUnwrap(member, grant, VaultA, 2).ShouldBeNull();
    }

    [Fact]
    public void AGrantForAnotherMember_DoesNotOpen()
    {
        using var member = Key.Create(KeyAgreementAlgorithm.X25519);
        using var intruder = Key.Create(KeyAgreementAlgorithm.X25519);
        var vaultKey = VaultKeys.Create();

        var grant = VaultKeys.WrapTo(
            vaultKey, member.PublicKey.Export(KeyBlobFormat.RawPublicKey), VaultA, 1);

        VaultKeys.TryUnwrap(intruder, grant, VaultA, 1).ShouldBeNull();
    }

    [Fact]
    public void AGrantContainingSomethingOtherThanAKey_IsRejected()
    {
        // A malicious granter can seal anything; the recipient finds out here rather than later, when
        // an unexplained item decryption failure would look like data corruption.
        using var member = Key.Create(KeyAgreementAlgorithm.X25519);

        var bogus = DshCrypto.SealTo(
            member.PublicKey.Export(KeyBlobFormat.RawPublicKey),
            new byte[16],
            DshAad.VaultKeyGrant(VaultA, 1));

        VaultKeys.TryUnwrap(member, bogus, VaultA, 1).ShouldBeNull();
    }

    [Fact]
    public void WrappingAKeyOfTheWrongLength_Throws()
    {
        using var member = Key.Create(KeyAgreementAlgorithm.X25519);

        Should.Throw<ArgumentException>(() => VaultKeys.WrapTo(
            new byte[16], member.PublicKey.Export(KeyBlobFormat.RawPublicKey), VaultA, 1));
    }

    // ---- Item data keys ----

    [Fact]
    public void ADataKey_RoundTripsUnderTheVaultKey()
    {
        var vaultKey = VaultKeys.Create();
        var dataKey = ItemKeys.CreateDataKey();

        var wrapped = ItemKeys.WrapDataKey(
            dataKey, vaultKey, CryptoSpec.AadResourceType.Host, ItemA, 1, 1);

        ItemKeys.TryUnwrapDataKey(
            vaultKey, wrapped, CryptoSpec.AadResourceType.Host, ItemA, 1, 1).ShouldBe(dataKey);
    }

    [Fact]
    public void ADataKeyFromAnotherItem_DoesNotOpen()
    {
        var vaultKey = VaultKeys.Create();
        var dataKey = ItemKeys.CreateDataKey();

        var wrapped = ItemKeys.WrapDataKey(
            dataKey, vaultKey, CryptoSpec.AadResourceType.Host, ItemA, 1, 1);

        ItemKeys.TryUnwrapDataKey(
            vaultKey, wrapped, CryptoSpec.AadResourceType.Host, ItemB, 1, 1).ShouldBeNull();
    }

    [Fact]
    public void ADataKeyFromAnotherItemVersion_DoesNotOpen()
    {
        var vaultKey = VaultKeys.Create();
        var dataKey = ItemKeys.CreateDataKey();

        var wrapped = ItemKeys.WrapDataKey(
            dataKey, vaultKey, CryptoSpec.AadResourceType.Host, ItemA, 1, 2);

        ItemKeys.TryUnwrapDataKey(
            vaultKey, wrapped, CryptoSpec.AadResourceType.Host, ItemA, 1, 1).ShouldBeNull();
    }

    [Fact]
    public void ADataKeyForAnotherResourceType_DoesNotOpen()
    {
        // An id collision across tables would otherwise let a credential's key open a host's.
        var vaultKey = VaultKeys.Create();
        var dataKey = ItemKeys.CreateDataKey();

        var wrapped = ItemKeys.WrapDataKey(
            dataKey, vaultKey, CryptoSpec.AadResourceType.Host, ItemA, 1, 1);

        ItemKeys.TryUnwrapDataKey(
            vaultKey, wrapped, CryptoSpec.AadResourceType.Credential, ItemA, 1, 1).ShouldBeNull();
    }

    // ---- Item payloads ----

    [Fact]
    public void APayload_RoundTripsUnderItsDataKey()
    {
        var dataKey = ItemKeys.CreateDataKey();
        var plaintext = "ssh -p 2222 dodo@bastion.internal"u8.ToArray();

        var envelope = ItemKeys.SealPayload(
            dataKey, plaintext, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1);

        ItemKeys.TryOpenPayload(
            dataKey, envelope, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1)
            .ShouldBe(plaintext);
    }

    [Fact]
    public void APayloadPastedOntoAnotherItem_DoesNotOpen()
    {
        // The headline property. A server holding every ciphertext cannot move one row's bytes onto
        // another row, which is a guarantee no amount of access control provides.
        var dataKey = ItemKeys.CreateDataKey();
        var plaintext = "secret"u8.ToArray();

        var envelope = ItemKeys.SealPayload(
            dataKey, plaintext, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1);

        ItemKeys.TryOpenPayload(
            dataKey, envelope, CryptoSpec.AadResourceType.Host, ItemB, DataKeyId, 1, 1)
            .ShouldBeNull();
    }

    [Fact]
    public void APayloadRolledBackToAnEarlierVersion_DoesNotOpen()
    {
        var dataKey = ItemKeys.CreateDataKey();

        var envelope = ItemKeys.SealPayload(
            dataKey, "v2"u8.ToArray(), CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 2);

        ItemKeys.TryOpenPayload(
            dataKey, envelope, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1)
            .ShouldBeNull();
    }

    [Fact]
    public void APayloadFromASupersededKeyGeneration_DoesNotOpen()
    {
        var dataKey = ItemKeys.CreateDataKey();

        var envelope = ItemKeys.SealPayload(
            dataKey, "old"u8.ToArray(), CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1);

        ItemKeys.TryOpenPayload(
            dataKey, envelope, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 2, 1)
            .ShouldBeNull();
    }

    [Fact]
    public void AMetadataBlob_DoesNotOpenAsAPayload()
    {
        // Distinct purposes on the same item, so a server cannot swap the two — a substitution that
        // would leave a client decrypting a password where it expected a display name.
        var dataKey = ItemKeys.CreateDataKey();

        var metadata = DshCrypto.Seal(
            dataKey,
            "display name"u8,
            DshAad.ItemMetadata(CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1));

        ItemKeys.TryOpenPayload(
            dataKey, metadata, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1)
            .ShouldBeNull();
    }

    [Fact]
    public void APayloadUnderAnotherDataKeyId_DoesNotOpen()
    {
        // content_key_id is part of the binding, which is what lets M5 add per-item grants without
        // reworking the AAD.
        var dataKey = ItemKeys.CreateDataKey();

        var envelope = ItemKeys.SealPayload(
            dataKey, "x"u8.ToArray(), CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1);

        ItemKeys.TryOpenPayload(
            dataKey, envelope, CryptoSpec.AadResourceType.Host, ItemA, Guid.Empty, 1, 1)
            .ShouldBeNull();
    }

    [Fact]
    public void TheFullChain_WorksEndToEnd()
    {
        // Enrollment through to a stored item: identity key, vault key sealed to it, data key wrapped
        // under the vault key, payload under the data key. Every link uses only what the layer above
        // it hands over, which is the shape the client actually follows.
        using var bundle = UserSecretBundle.Create(
            DateTimeOffset.FromUnixTimeMilliseconds(1_750_000_000_000));

        var vaultKey = VaultKeys.Create();
        var grant = VaultKeys.WrapTo(vaultKey, bundle.EncryptionPublicKey, VaultA, 1);

        var recoveredVaultKey = VaultKeys.TryUnwrap(bundle.EncryptionKey, grant, VaultA, 1);
        recoveredVaultKey.ShouldNotBeNull();

        var dataKey = ItemKeys.CreateDataKey();
        var wrappedDataKey = ItemKeys.WrapDataKey(
            dataKey, recoveredVaultKey, CryptoSpec.AadResourceType.Host, ItemA, 1, 1);

        var plaintext = """{"hostname":"bastion.internal","username":"dodo"}"""u8.ToArray();
        var payload = ItemKeys.SealPayload(
            dataKey, plaintext, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1);

        // Now the read path, holding nothing but the identity key and the stored ciphertexts.
        var readVaultKey = VaultKeys.TryUnwrap(bundle.EncryptionKey, grant, VaultA, 1);
        readVaultKey.ShouldNotBeNull();

        var readDataKey = ItemKeys.TryUnwrapDataKey(
            readVaultKey, wrappedDataKey, CryptoSpec.AadResourceType.Host, ItemA, 1, 1);
        readDataKey.ShouldNotBeNull();

        ItemKeys.TryOpenPayload(
            readDataKey, payload, CryptoSpec.AadResourceType.Host, ItemA, DataKeyId, 1, 1)
            .ShouldBe(plaintext);
    }
}