package keygen

import (
	"b612.me/starcrypto"
	"b612.me/starlog"
	"b612.me/staros"
	"crypto/ecdsa"
	"crypto/rsa"
	"github.com/spf13/cobra"
	"os"
	"time"
)

var k KeyGen
var startdate string
var duration int

var secret string
var path string
var key string
var outpath string

var sshPub bool

func init() {
	Cmd.Flags().StringVarP(&k.Type, "type", "t", "rsa", "Key Type: rsa, ecdsa, ed25519")
	Cmd.Flags().StringVarP(&k.Encrypt, "encrypt", "e", "", "Encrypt Key with Password (not recommended)")
	Cmd.Flags().IntVarP(&k.Bits, "bits", "b", 2048, "Key Bits Rsa: 1024, 2048, 4096 Ecdsa: 224, 256, 384, 521")
	Cmd.Flags().StringVarP(&k.Prefix, "prefix", "p", "mykey", "Output File Prefix")
	Cmd.Flags().StringVarP(&k.Outfolder, "outfolder", "o", ".", "Output Folder")
	Cmd.Flags().BoolVarP(&k.Force, "force", "f", false, "Force Overwrite")
	Cmd.Flags().StringVarP(&k.Country, "country", "c", "CN", "Country")
	Cmd.Flags().StringVarP(&k.Locality, "locality", "l", "Beijing", "Locality")
	Cmd.Flags().StringVarP(&k.Organization, "organization", "O", "B612", "Organization")
	Cmd.Flags().StringVarP(&k.OrganizationalUnit, "organizationalunit", "U", "B612", "OrganizationalUnit")
	Cmd.Flags().StringVarP(&k.CommonName, "commonname", "C", "Little Prince", "CommonName")
	Cmd.Flags().StringVarP(&startdate, "startdate", "s", "", "Cert Start Date")
	Cmd.Flags().IntVarP(&duration, "duration", "d", 3650, "Cert Duration")

	CmdEn.Flags().StringVarP(&secret, "secret", "s", "", "new Private Key Password,if empty,disable password")
	CmdEn.Flags().StringVarP(&path, "path", "p", "", "private key file path")
	CmdEn.Flags().StringVarP(&key, "key", "k", "", "private key old password,if empty,disable password")
	CmdEn.Flags().StringVarP(&outpath, "outpath", "o", "./newkey", "new key file output path")

	Cmd.AddCommand(CmdEn)

	CmdPub.Flags().StringVarP(&path, "path", "p", "", "private key file path")
	CmdPub.Flags().StringVarP(&outpath, "outpath", "o", "./public.key", "public key file output path")
	CmdPub.Flags().BoolVarP(&sshPub, "ssh", "s", false, "output ssh public key")
	Cmd.AddCommand(CmdPub)
}

var Cmd = &cobra.Command{
	Use:   "keygen",
	Short: "rsa与ecdsa密钥生成工具",
	Long:  "rsa与ecdsa密钥生成工具，支持加密私钥，生成证书",
	Run: func(cmd *cobra.Command, args []string) {
		var err error
		if startdate != "" {
			k.StartDate, err = time.Parse("2006-01-02", startdate)
			if err != nil {
				starlog.Errorln(err)
				os.Exit(1)
			}
		} else {
			k.StartDate = time.Now()
		}
		if duration < 0 {
			starlog.Errorln("duration should be positive")
			os.Exit(1)

		}
		k.EndDate = k.StartDate.AddDate(0, 0, duration)
		err = k.Gen()
		if err != nil {
			starlog.Errorln(err)
			os.Exit(1)
		}
		starlog.Infoln("Key Generated,Ouput to", k.Outfolder)
	},
}

var CmdEn = &cobra.Command{
	Use:   "pwd",
	Short: "encrypt/change private key",
	Run: func(cmd *cobra.Command, args []string) {
		if !staros.Exists(path) {
			starlog.Errorln("file not exists")
			os.Exit(1)
		}
		data, err := os.ReadFile(path)
		if err != nil {
			starlog.Errorln("read file error:", err)
			os.Exit(1)
		}
		priv, err := starcrypto.DecodePrivateKey(data, key)
		if err != nil {
			starlog.Errorln("decode private key error:", err)
			os.Exit(1)
		}
		data, err = starcrypto.EncodePrivateKey(priv, secret)
		if err != nil {
			starlog.Errorln("encode private key error:", err)
			os.Exit(1)
		}
		err = os.WriteFile(outpath, data, 0644)
		if err != nil {
			starlog.Errorln("write new file error:", err)
			os.Exit(1)
		}
		starlog.Infoln("new key saved to", outpath)
	},
}

var CmdPub = &cobra.Command{
	Use:   "pub",
	Short: "通过私钥生成公钥",
	Run: func(cmd *cobra.Command, args []string) {
		var pub any
		if !staros.Exists(path) {
			starlog.Errorln("file not exists")
			os.Exit(1)
		}
		data, err := os.ReadFile(path)
		if err != nil {
			starlog.Errorln("read file error:", err)
			os.Exit(1)
		}
		priv, err := starcrypto.DecodePrivateKey(data, key)
		if err != nil {
			starlog.Errorln("decode private key error:", err)
			os.Exit(1)
		}
		switch n := priv.(type) {
		case *rsa.PrivateKey:
			starlog.Infoln("found rsa private key")
			pub = n.Public()
		case *ecdsa.PrivateKey:
			starlog.Infoln("found ecdsa private key")
			pub = n.Public()
		default:
			starlog.Errorln("unknown private key type")
			os.Exit(1)
		}
		if sshPub {
			data, err = starcrypto.EncodeSSHPublicKey(pub)
			if err != nil {
				starlog.Errorln("encode ssh public key error:", err)
				os.Exit(1)
			}
		} else {
			data, err = starcrypto.EncodePublicKey(pub)
			if err != nil {
				starlog.Errorln("encode public key error:", err)
				os.Exit(1)
			}
		}
		starlog.Infoln("public key:", string(data))
		err = os.WriteFile(outpath, data, 0644)
		if err != nil {
			starlog.Errorln("write public key error:", err)
			os.Exit(1)
		}
		starlog.Infoln("public key saved to", outpath)
	},
}