package starcrypto

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

func EncodePrivateKey(private crypto.PrivateKey, secret string) ([]byte, error) {
	switch private.(type) {
	case *rsa.PrivateKey:
		return EncodeRsaPrivateKey(private.(*rsa.PrivateKey), secret)
	case *ecdsa.PrivateKey:
		return EncodeEcdsaPrivateKey(private.(*ecdsa.PrivateKey), secret)
	default:
		b, err := x509.MarshalPKCS8PrivateKey(private)
		if err != nil {
			return nil, err
		}
		if secret == "" {
			return pem.EncodeToMemory(&pem.Block{
				Bytes: b,
				Type:  "PRIVATE KEY",
			}), err
		}
		chiper := x509.PEMCipherAES256
		blk, err := x509.EncryptPEMBlock(rand.Reader, "PRIVATE KEY", b, []byte(secret), chiper)
		if err != nil {
			return nil, err
		}
		return pem.EncodeToMemory(blk), err
	}
}

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

func DecodePrivateKey(private []byte, password string) (crypto.PrivateKey, error) {
	blk, _ := pem.Decode(private)
	if blk == nil {
		return nil, errors.New("private key error")
	}
	switch blk.Type {
	case "RSA PRIVATE KEY":
		return DecodeRsaPrivateKey(private, password)
	case "EC PRIVATE KEY":
		return DecodeEcdsaPrivateKey(private, password)
	case "PRIVATE KEY":
		var prikey crypto.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.ParsePKCS8PrivateKey(bytes)
		if err != nil {
			return nil, err
		}
		return prikey, err
	default:
		return nil, errors.New("private key type error")
	}
}

func DecodePublicKey(pubStr []byte) (crypto.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, nil
}

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

func DecodeSSHPublicKey(pubStr []byte) (crypto.PublicKey, error) {
	return ssh.ParsePublicKey(pubStr)
}