// 提供RSA不对称加解密功能。
package rsa

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha512"
	"errors"
)

type KeyLength int
type SignAlgorithm int

const (
	RSA1024 KeyLength = 1024
	RSA2048 KeyLength = 2048
)

const (
	PSSSign SignAlgorithm = iota
	PKCS1Sign
)

type KeyPair struct {
	PublicKey  *rsa.PublicKey
	PrivateKey *rsa.PrivateKey
}

// 生成一个新的RSA密钥对。
func NewKeyPair(length KeyLength) (*KeyPair, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, int(length))
	if err != nil {
		return nil, &UnableToGenerateKeyPairError{err: err}
	}
	publicKey := privateKey.Public().(*rsa.PublicKey)
	return &KeyPair{
		PublicKey:  publicKey,
		PrivateKey: privateKey,
	}, nil
}

// 从一个PEM格式的密钥对中读取一套RSA密钥对。
// 如果获取到的公钥与私钥不配对，那么将自动使用私钥生成一个配对的公钥，以保证密钥对的完整性。
// 如果提供的证书中只存在公钥，那么将会保持私钥为空。
func NewFromPEM(cert []byte) (*KeyPair, error) {
	publicKey, err := DecodePublicKey(cert)
	if err != nil {
		if errors.Is(err, &KeyNotFoundError{}) {
			publicKey = nil
		} else {
			return nil, err
		}
	}
	privateKey, err := DecodePrivateKey(cert)
	if err != nil {
		if errors.Is(err, &KeyNotFoundError{}) {
			privateKey = nil
		} else {
			return nil, err
		}
	}
	if privateKey != nil {
		pubKeyFromPriKey := privateKey.Public().(*rsa.PublicKey)
		if publicKey == nil || !publicKey.Equal(pubKeyFromPriKey) {
			publicKey = pubKeyFromPriKey
		}
	}
	return &KeyPair{
		PublicKey:  publicKey,
		PrivateKey: privateKey,
	}, nil
}

// 加载一个新的公钥，将会自动清空私钥。
func (kp *KeyPair) LoadPublicKey(cert []byte) error {
	publicKey, err := DecodePublicKey(cert)
	if err != nil {
		return err
	}
	kp.PublicKey = publicKey
	kp.PrivateKey = nil
	return nil
}

// 加载一个私钥，将会自动生成一个配对的公钥。
func (kp *KeyPair) LoadPrivateKey(cert []byte) error {
	privateKey, err := DecodePrivateKey(cert)
	if err != nil {
		return err
	}
	kp.PrivateKey = privateKey
	kp.PublicKey = privateKey.Public().(*rsa.PublicKey)
	return nil
}

// 使用密钥对加密给定的数据，如果密钥对中不存在公钥则会返回错误。
func (kp KeyPair) Encrypt(data []byte, label ...[]byte) ([]byte, error) {
	if kp.PublicKey == nil {
		return nil, &AbsentPublicKeyError{}
	}
	var cipherLabel []byte
	if len(label) > 0 {
		cipherLabel = label[0]
	} else {
		cipherLabel = nil
	}
	hasher := sha512.New()
	cipherText, err := rsa.EncryptOAEP(hasher, rand.Reader, kp.PublicKey, data, cipherLabel)
	if err != nil {
		return nil, &EncryptionError{err: err}
	}
	return cipherText, nil
}

// 使用密钥对解密给定的数据，如果密钥对中不存在私钥则会返回错误。
func (kp KeyPair) Decrypt(data []byte, label ...[]byte) ([]byte, error) {
	if kp.PrivateKey == nil {
		return nil, &AbsentPrivateKeyError{}
	}
	var cipherLabel []byte
	if len(label) > 0 {
		cipherLabel = label[0]
	} else {
		cipherLabel = nil
	}
	hasher := sha512.New()
	plainText, err := rsa.DecryptOAEP(hasher, rand.Reader, kp.PrivateKey, data, cipherLabel)
	if err != nil {
		return nil, &DecryptionError{err: err}
	}
	return plainText, nil
}

// 对给定的数据使用私钥进行签名。
func (kp KeyPair) Sign(data []byte, algorithm ...SignAlgorithm) ([]byte, error) {
	if kp.PrivateKey == nil {
		return nil, &AbsentPrivateKeyError{}
	}
	hashedData := sha512.Sum512(data)
	switch append(algorithm, PSSSign)[0] {
	case PSSSign:
		signature, err := rsa.SignPSS(rand.Reader, kp.PrivateKey, crypto.SHA512, hashedData[:], &rsa.PSSOptions{SaltLength: 64})
		if err != nil {
			return nil, &SignError{err: err}
		}
		return signature, nil
	case PKCS1Sign:
		signature, err := rsa.SignPKCS1v15(rand.Reader, kp.PrivateKey, crypto.SHA512, hashedData[:])
		if err != nil {
			return nil, &SignError{err: err}
		}
		return signature, nil
	default:
		return nil, &InvalidSignMethodError{}
	}
}

// 对给定的数据和签名数据使用公钥进行验证。不返回任何错误信息即表示验证成功。
func (kp KeyPair) Verify(data, signature []byte, algorithm ...SignAlgorithm) error {
	if kp.PublicKey == nil {
		return &AbsentPublicKeyError{}
	}
	hashedData := sha512.Sum512(data)
	switch append(algorithm, PSSSign)[0] {
	case PSSSign:
		err := rsa.VerifyPSS(kp.PublicKey, crypto.SHA512, hashedData[:], signature, &rsa.PSSOptions{SaltLength: 64})
		if err != nil {
			return &VerifyError{err: err}
		}
		return nil
	case PKCS1Sign:
		err := rsa.VerifyPKCS1v15(kp.PublicKey, crypto.SHA512, hashedData[:], signature)
		if err != nil {
			return &VerifyError{err: err}
		}
		return nil
	default:
		return &InvalidSignMethodError{}
	}
}