RSA

---

package encryption_utils

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
)

// RSA struct to hold keys
type RSA struct {
	privateKey *rsa.PrivateKey
	publicKey  *rsa.PublicKey
}

// Initialize RSA with given keys
func NewRSA(privateKey, publicKey string) *RSA {
	rsa := &RSA{}
	prKey, _ := loadPrivateKey(privateKey)
	if prKey != nil {
		rsa.privateKey = prKey
	}

	pbKey, _ := loadPublicKey(publicKey)
	if pbKey != nil {
		rsa.publicKey = pbKey
	}

	return rsa
}

func (r *RSA) EncryptUsingPublicKey(data string) (encryptedData string, err error) {
	if r.publicKey == nil {
		return "", fmt.Errorf("public key is not set")
	}

	hash := sha256.New()
	encryptedBytes, err := rsa.EncryptOAEP(hash, rand.Reader, r.publicKey, []byte(data), nil)
	if err != nil {
		return "", fmt.Errorf("failed to encrypt data using OAEP: %v", err)
	}

	encryptedData = base64.StdEncoding.EncodeToString(encryptedBytes)

	return encryptedData, nil
}

func (r *RSA) DecryptUsingPrivateKey(encryptedData string) (decryptedData string, err error) {
	if r.privateKey == nil {
		return "", fmt.Errorf("private key is not set")
	}

	encryptedBytes, err := base64.StdEncoding.DecodeString(encryptedData)
	if err != nil {
		return "", fmt.Errorf("failed to decode base64 encoded data: %v", err)
	}

	hash := sha256.New()
	decryptedBytes, err := rsa.DecryptOAEP(hash, rand.Reader, r.privateKey, encryptedBytes, nil)
	if err != nil {
		return "", fmt.Errorf("failed to decrypt data using OAEP: %v", err)
	}

	decryptedData = string(decryptedBytes)

	return decryptedData, nil
}

// Helper functions to load keys from PEM encoded strings
func loadPrivateKey(privateKey string) (*rsa.PrivateKey, error) {
	pemPrivateKey := "-----BEGIN PRIVATE KEY-----\n" + privateKey + "\n-----END PRIVATE KEY-----"

	block, _ := pem.Decode([]byte(pemPrivateKey))
	if block == nil || block.Type != "PRIVATE KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing private key")
	}

	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	rsaKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA private key")
	}
	return rsaKey, nil
}

// Helper function to load public key from PEM encoded string
func loadPublicKey(publicKey string) (*rsa.PublicKey, error) {
	pemPublicKey := "-----BEGIN PUBLIC KEY-----\n" + publicKey + "\n-----END PUBLIC KEY-----"
	block, _ := pem.Decode([]byte(pemPublicKey))
	if block == nil || block.Type != "PUBLIC KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing public key")
	}

	key, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	rsaKey, ok := key.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA public key")
	}
	return rsaKey, nil
}