ModuleLatticeEndpoint.cs

using System;
using System.IO;
using Waher.Security.PQC;
 
namespace Waher.Networking.XMPP.P2P.E2E
{
    public abstract class ModuleLatticeEndpoint : E2eEndpoint
    {
        private const string preHashAlgorithm = "SHAKE-256";
 
        private readonly ML_KEM keyEncapsulationMechanism;
        private readonly ML_DSA signatureAlgorithm;
        private readonly ML_KEM_Keys keyEncapsulationMechanismKeys;
        private readonly ML_DSA_Keys signatureAlgorithmKeys;
        private readonly bool hasPrivateKey;
        private readonly byte[] publicKey;
        private readonly string publicKeyBase64;
 
        public ModuleLatticeEndpoint(ML_KEM KeyEncapsulationMechanism,
            ML_KEM_Keys KeyEncapsulationMechanismKeys, ML_DSA SignatureAlgorithm,
            ML_DSA_Keys SignatureAlgorithmKeys, IE2eSymmetricCipher DefaultSymmetricCipher)
            : base(DefaultSymmetricCipher)
        {
            this.keyEncapsulationMechanism = KeyEncapsulationMechanism;
            this.keyEncapsulationMechanismKeys = KeyEncapsulationMechanismKeys;
            this.signatureAlgorithm = SignatureAlgorithm;
            this.signatureAlgorithmKeys = SignatureAlgorithmKeys;
 
            this.hasPrivateKey = this.keyEncapsulationMechanismKeys.HasDecapsulationKey &&
                this.signatureAlgorithmKeys.HasPrivateKey;
 
            if (this.keyEncapsulationMechanismKeys.EncapsulationKey is null ||
                this.signatureAlgorithmKeys.PublicKey is null)
            {
                this.publicKey = null;
                this.publicKeyBase64 = null;
            }
            else
            {
                this.publicKey = new byte[this.keyEncapsulationMechanism.PublicKeyLength +
                    this.signatureAlgorithm.PublicKeyLength];
 
                Buffer.BlockCopy(this.keyEncapsulationMechanismKeys.EncapsulationKey, 0,
                    this.publicKey, 0, this.keyEncapsulationMechanism.PublicKeyLength);
 
                Buffer.BlockCopy(this.signatureAlgorithmKeys.PublicKey, 0,
                    this.publicKey, this.keyEncapsulationMechanism.PublicKeyLength,
                    this.signatureAlgorithm.PublicKeyLength);
 
                this.publicKeyBase64 = Convert.ToBase64String(this.publicKey);
            }
        }
 
        public ModuleLatticeEndpoint(byte[] PublicKey, ML_KEM KeyEncapsulationMechanism,
            ML_DSA SignatureAlgorithm, IE2eSymmetricCipher DefaultSymmetricCipher)
            : base(DefaultSymmetricCipher)
        {
            if (PublicKey is null)
                throw new ArgumentNullException(nameof(PublicKey));
 
            if (PublicKey.Length != KeyEncapsulationMechanism.PublicKeyLength +
                SignatureAlgorithm.PublicKeyLength)
            {
                throw new ArgumentException("Public key length does not match expected length for the specified key encapsulation mechanism and signature algorithm.", nameof(PublicKey));
            }
 
            this.keyEncapsulationMechanism = KeyEncapsulationMechanism;
            this.signatureAlgorithm = SignatureAlgorithm;
            this.hasPrivateKey = false;
            this.publicKey = PublicKey;
            this.publicKeyBase64 = Convert.ToBase64String(PublicKey);
 
            byte[] EncapsulationKey = new byte[this.keyEncapsulationMechanism.PublicKeyLength];
            byte[] SignaturePublicKey = new byte[this.signatureAlgorithm.PublicKeyLength];
 
            Buffer.BlockCopy(PublicKey, 0, EncapsulationKey, 0, EncapsulationKey.Length);
            Buffer.BlockCopy(PublicKey, EncapsulationKey.Length, SignaturePublicKey, 0, SignaturePublicKey.Length);
 
            this.keyEncapsulationMechanismKeys = ML_KEM_Keys.FromEncapsulationKey(EncapsulationKey);
            this.signatureAlgorithmKeys = ML_DSA_Keys.FromPublicKey(SignaturePublicKey);
        }
 
        public bool HasPrivateKey => this.hasPrivateKey;
 
        public override byte[] PublicKey => this.publicKey;
 
        public override string PublicKeyBase64 => this.publicKeyBase64;
 
        public ML_KEM KeyEncapsulationMechanism => this.keyEncapsulationMechanism;
 
        public ML_DSA SignatureAlgorithm => this.signatureAlgorithm;
 
        public override byte[] GetSharedSecretForEncryption(IE2eEndpoint RemoteEndpoint,
            IE2eSymmetricCipher Cipher, out byte[] CipherText)
        {
            if (!(RemoteEndpoint is ModuleLatticeEndpoint RemoteModuleLatticeEndpoint))
                throw new ArgumentException("Remote endpoint is not a Module Lattice endpoint.", nameof(RemoteEndpoint));
 
            if (!this.keyEncapsulationMechanismKeys.HasDecapsulationKey)
                throw new InvalidOperationException("Local endpoint does not have a private key for shared secret calculation.");
 
            if (this.keyEncapsulationMechanism.PublicKeyLength !=
                RemoteModuleLatticeEndpoint.keyEncapsulationMechanism.PublicKeyLength ||
                this.signatureAlgorithm.PublicKeyLength !=
                RemoteModuleLatticeEndpoint.signatureAlgorithm.PublicKeyLength)
            {
                throw new ArgumentException("Key lengths do not match.", nameof(RemoteEndpoint));
            }
 
            ML_KEM_Encapsulation Encapsulation = this.keyEncapsulationMechanism.Encapsulate(
                RemoteModuleLatticeEndpoint.keyEncapsulationMechanismKeys.EncapsulationKey);
 
            CipherText = Encapsulation.CipherText;
 
            return Encapsulation.SharedSecret;
        }
 
        public override byte[] GetSharedSecretForDecryption(IE2eEndpoint RemoteEndpoint,
            byte[] CipherText)
        {
            if (!(RemoteEndpoint is ModuleLatticeEndpoint RemoteModuleLatticeEndpoint))
                throw new ArgumentException("Remote endpoint is not a Module Lattice endpoint.", nameof(RemoteEndpoint));
 
            if (!this.keyEncapsulationMechanismKeys.HasDecapsulationKey)
                throw new InvalidOperationException("Local endpoint does not have a private key for shared secret calculation.");
 
            if (this.keyEncapsulationMechanism.PublicKeyLength !=
                RemoteModuleLatticeEndpoint.keyEncapsulationMechanism.PublicKeyLength ||
                this.signatureAlgorithm.PublicKeyLength !=
                RemoteModuleLatticeEndpoint.signatureAlgorithm.PublicKeyLength)
            {
                throw new ArgumentException("Key lengths do not match.", nameof(RemoteEndpoint));
            }
 
            return this.keyEncapsulationMechanism.Decapsulate(
                this.keyEncapsulationMechanismKeys.DecapsulationKey, CipherText);
        }
 
        public override bool SharedSecretUseCipherText => true;
 
        public override byte[] Sign(byte[] Data)
        {
            if (!this.hasPrivateKey)
                throw new InvalidOperationException("Signing requires private key.");
 
            byte[] Signature;
            bool Valid;
 
            do
            {
                Signature = this.signatureAlgorithm.Sign(this.signatureAlgorithmKeys,
                    Data, null, preHashAlgorithm, out _);
                Valid = this.signatureAlgorithm.Verify(this.signatureAlgorithmKeys, Data,
                    Signature, null, preHashAlgorithm);
            }
            while (!Valid);
 
            return Signature;
        }
 
        public override byte[] Sign(Stream Data)
        {
            if (!this.hasPrivateKey)
                throw new InvalidOperationException("Signing requires private key.");
 
            byte[] Signature;
            bool Valid;
 
            do
            {
                Signature = this.signatureAlgorithm.Sign(this.signatureAlgorithmKeys,
                    Data, null, preHashAlgorithm, out _);
                Valid = this.signatureAlgorithm.Verify(this.signatureAlgorithmKeys, Data,
                    Signature, null, preHashAlgorithm);
            }
            while (!Valid);
 
            return Signature;
        }
 
        public bool Verify(byte[] Data, byte[] PublicKey, byte[] Signature)
        {
            return this.signatureAlgorithm.Verify(ML_DSA_Keys.FromPublicKey(PublicKey), Data,
                Signature, null, preHashAlgorithm);
        }
 
        public bool Verify(Stream Data, byte[] PublicKey, byte[] Signature)
        {
            return this.signatureAlgorithm.Verify(ML_DSA_Keys.FromPublicKey(PublicKey), Data,
                Signature, null, preHashAlgorithm);
        }
 
        public override bool Verify(byte[] Data, byte[] Signature)
        {
            return this.signatureAlgorithm.Verify(this.signatureAlgorithmKeys, Data,
                Signature, null, preHashAlgorithm);
        }
 
        public override bool Verify(Stream Data, byte[] Signature)
        {
            return this.signatureAlgorithm.Verify(this.signatureAlgorithmKeys, Data,
                Signature, null, preHashAlgorithm);
        }
 
        public override bool Equals(object obj)
        {
            return obj is ModuleLatticeEndpoint MlObj &&
                this.publicKeyBase64 == MlObj.publicKeyBase64 &&
                this.hasPrivateKey == MlObj.hasPrivateKey;
        }
 
        public override int GetHashCode()
        {
            int Result = this.publicKeyBase64.GetHashCode();
            Result ^= Result << 5 ^ this.hasPrivateKey.GetHashCode();
 
            return Result;
        }
 
        protected void GeneratePrivateKeys(byte[] Secret, out ML_KEM_Keys KemKeys,
            out ML_DSA_Keys DsaKeys)
        {
            if (Secret is null)
                throw new ArgumentNullException(nameof(Secret));
 
            if (Secret.Length == 96)
            {
                byte[] SeedKem = new byte[64];
                byte[] SeedDsa = new byte[32];
 
                Buffer.BlockCopy(Secret, 0, SeedKem, 0, 64);
                Buffer.BlockCopy(Secret, 64, SeedDsa, 0, 32);
 
                KemKeys = this.keyEncapsulationMechanism.KeyGen_FromSeed(SeedKem, true);
                DsaKeys = this.signatureAlgorithm.KeyGen_Internal(SeedDsa, true);
 
                Array.Clear(SeedKem, 0, SeedKem.Length);
                Array.Clear(SeedDsa, 0, SeedDsa.Length);
            }
            else
                throw new ArgumentException("Invalid private key length.", nameof(Secret));
        }
 
        public byte[] ExportPrivateKey()
        {
            if (!(this.keyEncapsulationMechanismKeys?.HasDecapsulationKey ?? false) ||
                !(this.signatureAlgorithmKeys?.HasPrivateKey ?? false))
            {
                throw new InvalidOperationException("Endpoint has no private keys.");
            }
 
            if (this.keyEncapsulationMechanismKeys.Seed is null ||
                this.signatureAlgorithmKeys.Seed is null)
            {
                throw new InvalidOperationException("Endpoint has no seeds, which is required to export private key.");
            }
 
            byte[] KemKey = this.keyEncapsulationMechanismKeys.Seed;
            byte[] DsaKey = this.signatureAlgorithmKeys.Seed;
 
            int c = KemKey.Length;
            int d = DsaKey.Length;
 
            byte[] Result = new byte[c + d];
 
            Buffer.BlockCopy(KemKey, 0, Result, 0, c);
            Buffer.BlockCopy(DsaKey, 0, Result, c, d);
 
            return Result;
        }
    }
}