
// mlkem provides functions to encrypt and decrypt data using mlkem

package mlkem

// Public key is 1568 bytes long
// Private key is 64 bytes long
// Encrypted key is 1568 bytes long
// The key which is encrypted is 32 bytes long

// When the keys need to be encoded as strings:
// Mlkem Public keys are encoded Base64
// Mlkem private keys are encoded Hex

import "seekia/internal/encoding"

import cryptoMlkem "crypto/mlkem"

import "errors"

func VerifyMlkemPublicKeyString(inputKey string)bool{

	_, err := ReadMlkemPublicKeyString(inputKey)
	if (err != nil) {
		return false
	}

	return true
}

func ReadMlkemPublicKeyString(inputKey string)([1568]byte, error){

	decodedBytes, err := encoding.DecodeBase64StringToBytes(inputKey)
	if (err != nil) {
		return [1568]byte{}, errors.New("ReadMlkemPublicKeyString called with invalid key: " + inputKey)
	}

	if (len(decodedBytes) != 1568){
		return [1568]byte{}, errors.New("ReadMlkemPublicKeyString called with invalid key: " + inputKey)
	}

	// Now we verify that the public key is valid.

	_, err = cryptoMlkem.NewEncapsulationKey1024(decodedBytes)
	if (err != nil) {
		return [1568]byte{}, errors.New("ReadMlkemPublicKeyString called with invalid key: " + inputKey + ". Reason: " + err.Error())
	}

	publicKey := [1568]byte(decodedBytes)

	return publicKey, nil
}

func GetNewRandomPublicMlkemKey()([1568]byte, error){

	privateKeyObject, err := cryptoMlkem.GenerateKey1024()
	if (err != nil) { return [1568]byte{}, err }

	publicKeyObject := privateKeyObject.EncapsulationKey()
	publicKeyBytes := publicKeyObject.Bytes()

	publicKeyArray := [1568]byte(publicKeyBytes)

	return publicKeyArray, nil
}

func GetNewRandomPublicPrivateMlkemKeys()([1568]byte, [64]byte, error){

	privateKeyObject, err := cryptoMlkem.GenerateKey1024()
	if (err != nil) { return [1568]byte{}, [64]byte{}, err }

	publicKeyObject := privateKeyObject.EncapsulationKey()

	publicKeyBytes := publicKeyObject.Bytes()
	privateKeyBytes := privateKeyObject.Bytes()

	publicKeyArray := [1568]byte(publicKeyBytes)
	privateKeyArray := [64]byte(privateKeyBytes)

	return publicKeyArray, privateKeyArray, nil
}

//Outputs:
//	-[32]byte: New encrypted key
//	-[1568]byte: Mlkem encrypted key bytes
//	-error
func EncryptNewRandomKeyWithMlkem(publicMlkemKey [1568]byte)([32]byte, [1568]byte, error){	

	publicKeyObject, err := cryptoMlkem.NewEncapsulationKey1024(publicMlkemKey[:])
	if (err != nil) { return [32]byte{}, [1568]byte{}, err }

	newKey, encryptedResult := publicKeyObject.Encapsulate()

	newKeyArray := [32]byte(newKey)
	encryptedResultArray := [1568]byte(encryptedResult)

	return newKeyArray, encryptedResultArray, nil
}

//Outputs:
//	-bool: Able to decrypt
//	-[32]byte: Decrypted key
//	-error
func DecryptMlkemEncryptedKey(inputEncryptedKey [1568]byte, privateKey [64]byte)(bool, [32]byte, error){

	privateKeyObject, err := cryptoMlkem.NewDecapsulationKey1024(privateKey[:])
	if (err != nil) { return false, [32]byte{}, err }

	decryptedKey, err := privateKeyObject.Decapsulate(inputEncryptedKey[:])
	if (err != nil) {
		return false, [32]byte{}, nil
	}

	decryptedKeyArray := [32]byte(decryptedKey)

	return true, decryptedKeyArray, nil
}


