CSI-ES-2324/Projs/PD1/internal/utils/cryptoUtils/cryptoUtils.go

package cryptoUtils

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"log"
	"os"

	"golang.org/x/crypto/chacha20poly1305"
	"software.sslmate.com/src/go-pkcs12"
)

type KeyStore struct {
	cert        *x509.Certificate
	caCertChain []*x509.Certificate
	privKey     *rsa.PrivateKey
}

func (k KeyStore) GetCert() *x509.Certificate {
	return k.cert
}

func (k KeyStore) GetCACertChain() []*x509.Certificate {
	return k.caCertChain
}

func (k KeyStore) GetPrivKey() *rsa.PrivateKey {
	return k.privKey
}

func ExtractAllOIDValues(cert *x509.Certificate) map[string]string {
	oidValueMap := make(map[string]string)
	for _, name := range cert.Subject.Names {
		oid := name.Type.String()
		value := name.Value.(string)
		oidValueMap[oid] = value
	}
	return oidValueMap
}

func LoadKeyStore(keyStorePath string, password string) KeyStore {

	var privKey *rsa.PrivateKey

	certFile, err := os.ReadFile(keyStorePath)
	if err != nil {
		log.Panicln("Provided certificate couldn't be opened")
	}

	privKeyInterface, cert, caCerts, err := pkcs12.DecodeChain(certFile, password)
	if err != nil {
		log.Panicln("PKCS12 key store couldn't be decoded")
	}
	privKey, ok := privKeyInterface.(*rsa.PrivateKey)
	if !ok {
		log.Panicln("Failed to convert private key to RSA private key")
	}

	if err := privKey.Validate(); err != nil {
		log.Panicln("Private key is not valid")
	}
	return KeyStore{cert: cert, caCertChain: caCerts, privKey: privKey}
}

func (k *KeyStore) GetTLSConfig() *tls.Config {

	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: k.cert.Raw})
	privKeyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k.privKey)})
	certificate, err := tls.X509KeyPair(certPEM, privKeyPEM)
	if err != nil {
		log.Panicln("Could not load certificate and privkey to TLS", err)
	}

	config := &tls.Config{
		Certificates: []tls.Certificate{certificate},
	}
	return config
}

func (k *KeyStore) GetTLSConfigServer() *tls.Config {
	config := k.GetTLSConfig()

	//Add the CA certificate chain to a CertPool
	caCertPool := x509.NewCertPool()
	for _, caCert := range k.caCertChain {
		caCertPool.AddCert(caCert)
	}
	config.ClientCAs = caCertPool
	config.ClientAuth = tls.RequireAndVerifyClientCert

	return config
}

func (k *KeyStore) GetTLSConfigClient() *tls.Config {
	config := k.GetTLSConfig()

	//Add the CA certificate chain to a CertPool
	caCertPool := x509.NewCertPool()
	for _, caCert := range k.caCertChain {
		caCertPool.AddCert(caCert)
	}
	config.RootCAs = caCertPool

	//TODO: FIX THE VERIFICATION OF THE SERVER
	//config.ServerName = "SERVER"

	return config
}

func (k KeyStore) EncryptMessageContent(receiverCert *x509.Certificate, content []byte) []byte {
	// Digital envolope

	// Create a random symmetric key
	dataKey := make([]byte, 32)
	if _, err := rand.Read(dataKey); err != nil {
		log.Panicln("Could not create dataKey properly: ", err)
	}

	cipher, err := chacha20poly1305.New(dataKey)
	if err != nil {
		log.Panicln("Could not create cipher: ", err)
	}

	nonce := make([]byte, cipher.NonceSize(), cipher.NonceSize()+len(content)+cipher.Overhead())
	if _, err := rand.Read(nonce); err != nil {
		log.Panicln("Could not create data nonce properly: ", err)
	}

	// sign the message and append the signature
	hashedContent := sha256.Sum256(content)
	signature, err := rsa.SignPKCS1v15(nil, k.privKey, crypto.SHA256, hashedContent[:])
	if err != nil {
		log.Panicln("Could not create content signature: ", err)
	}
	content = pair(content, signature)
	ciphertext := cipher.Seal(nonce, nonce, content, nil)

	// crypto/rand.Reader is a good source of entropy for randomizing the
	// encryption function.
	receiverPubKey := receiverCert.PublicKey.(*rsa.PublicKey)
	encryptedDataKey, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, receiverPubKey, dataKey, nil)
	if err != nil {
		log.Panicln("Could not encrypt dataKey: ", err)
	}

	return pair(encryptedDataKey, ciphertext)
}

func (k KeyStore) DecryptMessageContent(senderCert *x509.Certificate, cipherContent []byte) []byte {

   return nil
}

func pair(l []byte, r []byte) []byte {
	length := len(l)
	lenBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(lenBytes, uint16(length))

	lWithLen := append(lenBytes, l...)
	return append(lWithLen, r...)
}

func unPair(pair []byte) ([]byte, []byte) {
	lenBytes := pair[:2]
	pair = pair[2:]
	length := binary.BigEndian.Uint16(lenBytes)
	l := pair[:length]
	r := pair[length:]
	return l, r
}