package starcrypto

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"golang.org/x/crypto/ssh"
	"math/big"
)

func GenerateRsaKey(bits int) (*rsa.PrivateKey, *rsa.PublicKey, error) {
	private, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, nil, err
	}
	return private, &private.PublicKey, nil

}

func EncodeRsaPrivateKey(private *rsa.PrivateKey, secret string) ([]byte, error) {
	if secret == "" {
		return pem.EncodeToMemory(&pem.Block{
			Bytes: x509.MarshalPKCS1PrivateKey(private),
			Type:  "RSA PRIVATE KEY",
		}), nil
	}
	chiper := x509.PEMCipherAES256
	blk, err := x509.EncryptPEMBlock(rand.Reader, "RSA PRIVATE KEY", x509.MarshalPKCS1PrivateKey(private), []byte(secret), chiper)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(blk), err
}

func EncodeRsaPublicKey(public *rsa.PublicKey) ([]byte, error) {
	publicBytes, err := x509.MarshalPKIXPublicKey(public)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{
		Bytes: publicBytes,
		Type:  "PUBLIC KEY",
	}), nil
}

func DecodeRsaPrivateKey(private []byte, password string) (*rsa.PrivateKey, error) {
	var prikey *rsa.PrivateKey
	var err error
	var bytes []byte
	blk, _ := pem.Decode(private)
	if blk == nil {
		return nil, errors.New("private key error！")
	}
	if password != "" {
		tmp, err := x509.DecryptPEMBlock(blk, []byte(password))
		if err != nil {
			return nil, err
		}
		bytes = tmp
	} else {
		bytes = blk.Bytes
	}
	prikey, err = x509.ParsePKCS1PrivateKey(bytes)
	if err != nil {
		tmp, err := x509.ParsePKCS8PrivateKey(bytes)
		if err != nil {
			return nil, err
		}
		prikey = tmp.(*rsa.PrivateKey)
	}
	return prikey, err
}

func DecodeRsaPublicKey(pubStr []byte) (*rsa.PublicKey, error) {
	blk, _ := pem.Decode(pubStr)
	if blk == nil {
		return nil, errors.New("public key error")
	}
	pub, err := x509.ParsePKIXPublicKey(blk.Bytes)
	if err != nil {
		return nil, err
	}
	return pub.(*rsa.PublicKey), nil
}

func EncodeRsaSSHPublicKey(public *rsa.PublicKey) ([]byte, error) {
	publicKey, err := ssh.NewPublicKey(public)
	if err != nil {
		return nil, err
	}
	return ssh.MarshalAuthorizedKey(publicKey), nil
}

func GenerateRsaSSHKeyPair(bits int, secret string) (string, string, error) {
	pkey, pubkey, err := GenerateRsaKey(bits)
	if err != nil {
		return "", "", err
	}

	pub, err := EncodeRsaSSHPublicKey(pubkey)
	if err != nil {
		return "", "", err
	}

	priv, err := EncodeRsaPrivateKey(pkey, secret)
	if err != nil {
		return "", "", err
	}
	return string(priv), string(pub), nil
}

// RSAEncrypt RSA公钥加密
func RSAEncrypt(pub *rsa.PublicKey, data []byte) ([]byte, error) {
	return rsa.EncryptPKCS1v15(rand.Reader, pub, data)
}

// RSADecrypt RSA私钥解密
func RSADecrypt(prikey *rsa.PrivateKey, data []byte) ([]byte, error) {
	return rsa.DecryptPKCS1v15(rand.Reader, prikey, data)
}

// RSASign RSA私钥签名加密
func RSASign(msg, priKey []byte, password string, hashType crypto.Hash) ([]byte, error) {
	var prikey *rsa.PrivateKey
	var err error
	var bytes []byte
	blk, _ := pem.Decode(priKey)
	if blk == nil {
		return []byte{}, errors.New("private key error！")
	}
	if password != "" {
		tmp, err := x509.DecryptPEMBlock(blk, []byte(password))
		if err != nil {
			return []byte{}, err
		}
		bytes = tmp
	} else {
		bytes = blk.Bytes
	}
	prikey, err = x509.ParsePKCS1PrivateKey(bytes)
	if err != nil {
		tmp, err := x509.ParsePKCS8PrivateKey(bytes)
		if err != nil {
			return []byte{}, err
		}
		prikey = tmp.(*rsa.PrivateKey)
	}
	hashMethod := hashType.New()
	_, err = hashMethod.Write(msg)
	if err != nil {
		return nil, err
	}
	return rsa.SignPKCS1v15(rand.Reader, prikey, hashType, hashMethod.Sum(nil))
}

// RSAVerify RSA公钥签名验证
func RSAVerify(data, msg, pubKey []byte, hashType crypto.Hash) error {
	blk, _ := pem.Decode(pubKey)
	if blk == nil {
		return errors.New("public key error！")
	}
	pubkey, err := x509.ParsePKIXPublicKey(blk.Bytes)
	if err != nil {
		return err
	}
	hashMethod := hashType.New()
	_, err = hashMethod.Write(msg)
	if err != nil {
		return err
	}
	return rsa.VerifyPKCS1v15(pubkey.(*rsa.PublicKey), hashType, hashMethod.Sum(nil), data)
}

// copy from crypt/rsa/pkcs1v5.go
var hashPrefixes = map[crypto.Hash][]byte{
	crypto.MD5:       {0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	crypto.SHA1:      {0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	crypto.SHA224:    {0x30, 0x2d, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04, 0x05, 0x00, 0x04, 0x1c},
	crypto.SHA256:    {0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	crypto.SHA384:    {0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	crypto.SHA512:    {0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	crypto.MD5SHA1:   {}, // A special TLS case which doesn't use an ASN1 prefix.
	crypto.RIPEMD160: {0x30, 0x20, 0x30, 0x08, 0x06, 0x06, 0x28, 0xcf, 0x06, 0x03, 0x00, 0x31, 0x04, 0x14},
}

// copy from crypt/rsa/pkcs1v5.go
func encrypt(c *big.Int, pub *rsa.PublicKey, m *big.Int) *big.Int {
	e := big.NewInt(int64(pub.E))
	c.Exp(m, e, pub.N)
	return c
}

// copy from crypt/rsa/pkcs1v5.go
func pkcs1v15HashInfo(hash crypto.Hash, inLen int) (hashLen int, prefix []byte, err error) {
	// Special case: crypto.Hash(0) is used to indicate that the data is
	// signed directly.
	if hash == 0 {
		return inLen, nil, nil
	}

	hashLen = hash.Size()
	if inLen != hashLen {
		return 0, nil, errors.New("crypto/rsa: input must be hashed message")
	}
	prefix, ok := hashPrefixes[hash]
	if !ok {
		return 0, nil, errors.New("crypto/rsa: unsupported hash function")
	}
	return
}

// copy from crypt/rsa/pkcs1v5.go
func leftPad(input []byte, size int) (out []byte) {
	n := len(input)
	if n > size {
		n = size
	}
	out = make([]byte, size)
	copy(out[len(out)-n:], input)
	return
}
func unLeftPad(input []byte) (out []byte) {
	n := len(input)
	t := 2
	for i := 2; i < n; i++ {
		if input[i] == 0xff {
			t = t + 1
		} else {
			if input[i] == input[0] {
				t = t + int(input[1])
			}
			break
		}
	}
	out = make([]byte, n-t)
	copy(out, input[t:])
	return
}

// copy&modified from crypt/rsa/pkcs1v5.go
func publicDecrypt(pub *rsa.PublicKey, hash crypto.Hash, hashed []byte, sig []byte) (out []byte, err error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return nil, err
	}

	tLen := len(prefix) + hashLen
	k := (pub.N.BitLen() + 7) / 8
	if k < tLen+11 {
		return nil, fmt.Errorf("length illegal")
	}

	c := new(big.Int).SetBytes(sig)
	m := encrypt(new(big.Int), pub, c)
	em := leftPad(m.Bytes(), k)
	out = unLeftPad(em)

	err = nil
	return
}

func RSAEncryptByPrivkey(privt *rsa.PrivateKey, data []byte) ([]byte, error) {
	signData, err := rsa.SignPKCS1v15(nil, privt, crypto.Hash(0), data)
	if err != nil {
		return nil, err
	}
	return signData, nil
}

func RSADecryptByPubkey(pub *rsa.PublicKey, data []byte) ([]byte, error) {
	decData, err := publicDecrypt(pub, crypto.Hash(0), nil, data)
	if err != nil {
		return nil, err
	}
	return decData, nil
}