package keygen

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"os"

	"github.com/spf13/cobra"
	"src.lwithers.me.uk/go/rsa/pkg/pemfile"
)

var (
	bits  int
	pkcs1 bool
)

// Register the "keygen" subcommand.
func Register(root *cobra.Command) {
	cmd := &cobra.Command{
		Use:   "keygen",
		Short: "Generate a new private key and save to file",
		Run:   Keygen,
		Args:  cobra.ExactArgs(1),
	}

	cmd.Flags().IntVarP(&bits, "bits", "b", 3072, "Key size in bits")
	cmd.Flags().BoolVarP(&pkcs1, "pkcs1", "", false, "Write key as PKCS#1 rather than PKCS#8")

	root.AddCommand(cmd)
}

// Keygen will generate a new RSA private key and save it to a file.
func Keygen(cmd *cobra.Command, args []string) {
	filename := args[0]

	key, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to generate new key (%d bits): %v\n", bits, err)
		os.Exit(1)
	}

	// prepare DER-form data
	var der []byte
	if pkcs1 {
		der = x509.MarshalPKCS1PrivateKey(key)
	} else {
		der, err = x509.MarshalPKCS8PrivateKey(key)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to marshal PKCS#8 private key: %v\n", err)
			os.Exit(1)
		}
	}

	// prepare PEM-form data
	typ := pemfile.TypePKCS8PrivateKey
	if pkcs1 {
		typ = pemfile.TypePKCS1PrivateKey
	}
	x509raw := pem.EncodeToMemory(&pem.Block{
		Type:  typ,
		Bytes: der,
	})

	// write to file
	if err := os.WriteFile(filename, x509raw, 0600); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}