package keygen

import (
	"b612.me/starcrypto"
	"b612.me/staros"
	"crypto"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io"
	"math/big"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type KeyGen struct {
	Type      string
	Encrypt   string
	Bits      int
	Outfolder string
	Prefix    string
	Force     bool
	//
	Country            string
	Locality           string
	Organization       string
	OrganizationalUnit string
	CommonName         string
	StartDate          time.Time
	EndDate            time.Time
}

func (k *KeyGen) Gen() error {
	if !k.Force && staros.Exists(filepath.Join(k.Outfolder, k.Prefix+".crt")) {
		return errors.New("crt file exists")
	}
	if !k.Force && staros.Exists(filepath.Join(k.Outfolder, k.Prefix)) {
		return errors.New("key file exists")
	}
	if !k.Force && staros.Exists(filepath.Join(k.Outfolder, k.Prefix+".pub")) {
		return errors.New("ssh pub file exists")
	}
	if !k.Force && staros.Exists(filepath.Join(k.Outfolder, k.Prefix+".openssh")) {
		return errors.New("ssh priv file exists")
	}
	if !k.Force && staros.Exists(filepath.Join(k.Outfolder, k.Prefix+".key.pub")) {
		return errors.New("pub file exists")
	}
	var sshPubByte, sshPrivByte, keyPubByte, keyPrivByte, Crt []byte

	var priv, pub any
	var err error
	switch strings.ToLower(k.Type) {
	case "rsa":
		priv, pub, err = starcrypto.GenerateRsaKey(k.Bits)
		if err != nil {
			return err
		}
	case "ecdsa", "ecdh":
		var cr elliptic.Curve
		switch k.Bits {
		case 224:
			cr = elliptic.P224()
		case 256:
			cr = elliptic.P256()
		case 384:
			cr = elliptic.P384()
		case 521:
			cr = elliptic.P521()
		default:
			return errors.New("invalid bits,should be 224,256,384,521")
		}
		priv, pub, err = starcrypto.GenerateEcdsaKey(cr)
		if err != nil {
			return err
		}
	case "ed25519":
		pub, priv, err = ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return err
		}
	default:
		return errors.New("invalid key type,only support rsa,ecdsa")
	}
	sshPubByte, err = starcrypto.EncodeSSHPublicKey(pub)
	if err != nil {
		return err
	}
	keyPubByte, err = starcrypto.EncodePublicKey(pub)
	if err != nil {
		return err
	}
	keyPrivByte, err = starcrypto.EncodePrivateKey(priv, k.Encrypt)
	if err != nil {
		return err
	}
	var block *pem.Block
	if k.Encrypt != "" {
		block, err = ssh.MarshalPrivateKey(priv, "")
	} else {
		block, err = ssh.MarshalPrivateKeyWithPassphrase(priv, "", []byte(k.Encrypt))
	}
	if err != nil {
		return err
	}
	sshPrivByte = pem.EncodeToMemory(block)
	if err != nil {
		return err
	}
	_, Crt, err = k.GenerateCert(priv)
	if err != nil {
		return err
	}
	err = os.WriteFile(filepath.Join(k.Outfolder, k.Prefix+".openssh"), sshPrivByte, 0644)
	if err != nil {
		return err
	}
	err = os.WriteFile(filepath.Join(k.Outfolder, k.Prefix+".crt"), Crt, 0644)
	if err != nil {
		return err
	}
	err = os.WriteFile(filepath.Join(k.Outfolder, k.Prefix), keyPrivByte, 0644)
	if err != nil {
		return err
	}
	err = os.WriteFile(filepath.Join(k.Outfolder, k.Prefix+".pub"), sshPubByte, 0644)
	if err != nil {
		return err
	}
	err = os.WriteFile(filepath.Join(k.Outfolder, k.Prefix+".key.pub"), keyPubByte, 0644)
	if err != nil {
		return err
	}
	return nil
}

func (k *KeyGen) GenerateCert(priv crypto.PrivateKey) ([]byte, []byte, error) {
	//csr,pub
	tmpByte := make([]byte, 64)
	io.ReadFull(rand.Reader, tmpByte)
	hexStr := starcrypto.String(tmpByte)
	data, _ := hex.DecodeString(hexStr)
	num := new(big.Int).SetBytes(data)
	var country, locality, organization, organizationalUnit []string
	if k.Country != "" {
		country = []string{k.Country}
	}
	if k.Locality != "" {
		locality = []string{k.Locality}
	}
	if k.Organization != "" {
		organization = []string{k.Organization}
	}
	if k.OrganizationalUnit != "" {
		organizationalUnit = []string{k.OrganizationalUnit}

	}
	var rootCsr = &x509.Certificate{
		Version:      3,
		SerialNumber: num,
		Subject: pkix.Name{
			Country:            country,
			Locality:           locality,
			Organization:       organization,
			OrganizationalUnit: organizationalUnit,
			CommonName:         k.CommonName,
		},
		NotBefore:             k.StartDate,
		NotAfter:              k.EndDate,
		BasicConstraintsValid: true,
		IsCA:                  false,
		MaxPathLenZero:        false,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageAny},
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign | x509.KeyUsageKeyEncipherment | x509.KeyUsageKeyAgreement | x509.KeyUsageDigitalSignature,
	}
	var cert []byte
	var err error
	switch priv.(type) {
	case *rsa.PrivateKey:
		cert, err = MakeCert(priv.(*rsa.PrivateKey), rootCsr, rootCsr, priv.(*rsa.PrivateKey).Public())
		if err != nil {
			return nil, nil, err
		}
	case *ecdsa.PrivateKey:
		cert, err = MakeCert(priv.(*ecdsa.PrivateKey), rootCsr, rootCsr, priv.(*ecdsa.PrivateKey).Public())
	case ed25519.PrivateKey:
		cert, err = MakeCert(priv.(ed25519.PrivateKey), rootCsr, rootCsr, priv.(ed25519.PrivateKey).Public())
	default:
		return nil, nil, errors.New("invalid private key type:" + fmt.Sprintf("%T", priv))
	}

	csrPem := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE REQUEST",
		Bytes: rootCsr.Raw,
	})
	return csrPem, cert, nil
}

func MakeCert(caKey any, caCrt *x509.Certificate, csr *x509.Certificate, pub any) ([]byte, error) {
	der, err := x509.CreateCertificate(rand.Reader, csr, caCrt, pub, caKey)
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		return nil, err
	}
	certBlock := &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert.Raw,
	}
	pemData := pem.EncodeToMemory(certBlock)
	return pemData, nil
}