package starcrypto

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"golang.org/x/crypto/ssh"
)

func GenerateEcdsaKey(pubkeyCurve elliptic.Curve) (*ecdsa.PrivateKey, *ecdsa.PublicKey, error) {
	// 随机挑选基点,生成私钥
	priv, err := ecdsa.GenerateKey(pubkeyCurve, rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return priv, &priv.PublicKey, nil

}

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

func EncodeEcdsaPublicKey(public *ecdsa.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 DecodeEcdsaPrivateKey(private []byte, password string) (*ecdsa.PrivateKey, error) {
	var prikey *ecdsa.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.ParseECPrivateKey(bytes)
	if err != nil {
		tmp, err := x509.ParsePKCS8PrivateKey(bytes)
		if err != nil {
			return nil, err
		}
		prikey = tmp.(*ecdsa.PrivateKey)
	}
	return prikey, err
}

func DecodeEcdsaPublicKey(pubStr []byte) (*ecdsa.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.(*ecdsa.PublicKey), nil
}

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

func GenerateEcdsaSSHKeyPair(pubkeyCurve elliptic.Curve, secret string) (string, string, error) {
	pkey, pubkey, err := GenerateEcdsaKey(pubkeyCurve)
	if err != nil {
		return "", "", err
	}

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

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