Envelope/Backend/Database/Seeder/MessageSeeder.go

package Seeder

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/pem"
	"fmt"
	"hash"

	"git.tovijaeschke.xyz/tovi/Envelope/Backend/Database"
	"git.tovijaeschke.xyz/tovi/Envelope/Backend/Models"
	"git.tovijaeschke.xyz/tovi/Envelope/Backend/Util"
	"github.com/gofrs/uuid"
)

// EncryptWithPublicKey encrypts data with public key
func encryptWithPublicKey(msg []byte, pub *rsa.PublicKey) []byte {
	var (
		hash hash.Hash
	)

	hash = sha256.New()
	ciphertext, err := rsa.EncryptOAEP(hash, rand.Reader, pub, msg, nil)
	if err != nil {
		panic(err)
	}
	return ciphertext
}

func PKCS5Padding(ciphertext []byte, blockSize int, after int) []byte {
	var (
		padding int
		padtext []byte
	)
	padding = (blockSize - len(ciphertext)%blockSize)
	padtext = bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func generateAesKey() (*pem.Block, cipher.BlockMode) {
	var (
		pemBlock *pem.Block
		block    cipher.Block

		bKey []byte
		bIV  []byte

		err error
	)

	bKey = make([]byte, 32)
	_, err = rand.Read(bKey)
	if err != nil {
		panic(err)
	}
	bIV = make([]byte, 16)
	_, err = rand.Read(bIV)
	if err != nil {
		panic(err)
	}

	pemBlock = &pem.Block{
		Type:  "AES KEY",
		Bytes: bKey,
	}

	block, err = aes.NewCipher(bKey)
	if err != nil {
		panic(err)
	}

	return pemBlock, cipher.NewCBCEncrypter(block, bIV)
}

func seedMessage(
	primaryUser Models.User,
	primaryUserThreadKey, secondaryUserThreadKey string,
	thread Models.MessageThread,
	i int,
) error {
	var (
		message     Models.Message
		messageData Models.MessageData

		messagePemBlock *pem.Block
		messageMode     cipher.BlockMode

		plaintext          string
		dataCiphertext     []byte
		senderIdCiphertext []byte

		bPlaintext         []byte
		bSenderIdPlaintext []byte

		err error
	)

	plaintext = "Test Message"
	bPlaintext = PKCS5Padding([]byte(plaintext), aes.BlockSize, len(plaintext))
	bSenderIdPlaintext = PKCS5Padding(primaryUser.ID.Bytes(), aes.BlockSize, len(primaryUser.ID.Bytes()))

	dataCiphertext = make([]byte, len(bPlaintext))
	senderIdCiphertext = make([]byte, len(bSenderIdPlaintext))

	messagePemBlock, messageMode = generateAesKey()

	messageMode.CryptBlocks(dataCiphertext, bPlaintext)
	messageMode.CryptBlocks(senderIdCiphertext, bSenderIdPlaintext)

	messageData = Models.MessageData{
		Data:     dataCiphertext,
		SenderID: senderIdCiphertext,
	}

	message = Models.Message{
		MessageData:      messageData,
		SymmetricKey:     encryptWithPublicKey(pem.EncodeToMemory(messagePemBlock), decodedPublicKey),
		MessageThreadKey: primaryUserThreadKey,
	}

	err = Database.CreateMessage(&message)
	if err != nil {
		return err
	}

	// The symmetric key would be encrypted with secondary users public key in production
	// But due to using the same pub/priv key pair for all users, we will just duplicate it
	message = Models.Message{
		MessageDataID:    message.MessageDataID,
		SymmetricKey:     encryptWithPublicKey(pem.EncodeToMemory(messagePemBlock), decodedPublicKey),
		MessageThreadKey: secondaryUserThreadKey,
	}

	err = Database.CreateMessage(&message)
	if err != nil {
		return err
	}

	return err
}

func seedMessageThread(threadPemBlock *pem.Block, threadMode cipher.BlockMode) (Models.MessageThread, error) {
	var (
		messageThread Models.MessageThread

		name           string
		bNamePlaintext []byte
		nameCiphertext []byte

		err error
	)

	name = "Test Conversation"

	bNamePlaintext = PKCS5Padding([]byte(name), aes.BlockSize, len(name))
	nameCiphertext = make([]byte, len(bNamePlaintext))

	threadMode.CryptBlocks(nameCiphertext, bNamePlaintext)

	messageThread = Models.MessageThread{
		Name: nameCiphertext,
	}

	err = Database.CreateMessageThread(&messageThread)
	return messageThread, err
}

func seedUpdateMessageThreadUsers(
	userJson string,
	threadPemBlock *pem.Block,
	threadMode cipher.BlockMode,
	messageThread Models.MessageThread,
) (Models.MessageThread, error) {
	var (
		bUsersPlaintext []byte
		usersCiphertext []byte
		err             error
	)

	bUsersPlaintext = PKCS5Padding([]byte(userJson), aes.BlockSize, len(userJson))
	usersCiphertext = make([]byte, len(bUsersPlaintext))

	threadMode.CryptBlocks(usersCiphertext, bUsersPlaintext)

	messageThread.Users = usersCiphertext
	err = Database.UpdateMessageThread(&messageThread)
	return messageThread, err
}

func seedMessageThreadUser(
	user Models.User,
	threadID uuid.UUID,
	messageThreadKey string,
	threadPemBlock *pem.Block,
	threadMode cipher.BlockMode,
) (Models.MessageThreadUser, error) {
	var (
		messageThreadUser Models.MessageThreadUser

		bThreadIdPlaintext []byte
		threadIdCiphertext []byte

		bKeyPlaintext []byte
		keyCiphertext []byte

		bAdminPlaintext []byte
		adminCiphertext []byte

		err error
	)

	bThreadIdPlaintext = PKCS5Padding(threadID.Bytes(), aes.BlockSize, len(threadID.String()))
	threadIdCiphertext = make([]byte, len(bThreadIdPlaintext))

	bKeyPlaintext = PKCS5Padding([]byte(messageThreadKey), aes.BlockSize, len(messageThreadKey))
	keyCiphertext = make([]byte, len(bKeyPlaintext))

	bAdminPlaintext = PKCS5Padding([]byte("true"), aes.BlockSize, len("true"))
	adminCiphertext = make([]byte, len(bAdminPlaintext))

	threadMode.CryptBlocks(threadIdCiphertext, bThreadIdPlaintext)
	threadMode.CryptBlocks(keyCiphertext, bKeyPlaintext)
	threadMode.CryptBlocks(adminCiphertext, bAdminPlaintext)

	messageThreadUser = Models.MessageThreadUser{
		UserID:           user.ID,
		MessageThreadID:  threadIdCiphertext,
		MessageThreadKey: keyCiphertext,
		Admin:            adminCiphertext,
		SymmetricKey:     encryptWithPublicKey(pem.EncodeToMemory(threadPemBlock), decodedPublicKey),
	}

	err = Database.CreateMessageThreadUser(&messageThreadUser)
	return messageThreadUser, err
}

func SeedMessages() {
	var (
		messageThread  Models.MessageThread
		threadPemBlock *pem.Block
		threadMode     cipher.BlockMode

		primaryUser            Models.User
		primaryUserThreadKey   string
		secondaryUser          Models.User
		secondaryUserThreadKey string

		userJson string

		thread Models.MessageThread
		i      int
		err    error
	)

	threadPemBlock, threadMode = generateAesKey()
	messageThread, err = seedMessageThread(threadPemBlock, threadMode)
	primaryUserThreadKey = Util.RandomString(32)
	secondaryUserThreadKey = Util.RandomString(32)

	primaryUser, err = Database.GetUserByUsername("testUser")
	if err != nil {
		panic(err)
	}

	_, err = seedMessageThreadUser(
		primaryUser,
		messageThread.ID,
		primaryUserThreadKey,
		threadPemBlock,
		threadMode,
	)

	secondaryUser, err = Database.GetUserByUsername("testUser2")
	if err != nil {
		panic(err)
	}

	_, err = seedMessageThreadUser(
		secondaryUser,
		messageThread.ID,
		secondaryUserThreadKey,
		threadPemBlock,
		threadMode,
	)

	userJson = fmt.Sprintf(
		`
[
	{
		"id": "%s",
		"username": "%s",
		"admin": "true"
	},
	{
		"id": "%s",
		"username": "%s",
		"admin": "true"
	}
]
		`,
		primaryUser.ID.String(),
		primaryUser.Username,
		secondaryUser.ID.String(),
		secondaryUser.Username,
	)

	messageThread, err = seedUpdateMessageThreadUsers(
		userJson,
		threadPemBlock,
		threadMode,
		messageThread,
	)

	for i = 0; i <= 20; i++ {
		err = seedMessage(
			primaryUser,
			primaryUserThreadKey,
			secondaryUserThreadKey,
			thread,
			i,
		)
		if err != nil {
			panic(err)
		}
	}
}