package security

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"golang.org/x/crypto/ssh"
	"log"
	"os"
	"path"
)

const (
	defaultPrivateKeyName = "id_rsa"
	defaultPublicKeyName  = "id_rsa.pub"
)

// GenerateNewRSAKey generates a new RSA key, writes it to a file and return the private key
func GenerateNewRSAKey(savePath string, bitSize uint16) (*rsa.PrivateKey, error) {
	pkPath := path.Join(savePath, defaultPrivateKeyName)
	pubPath := path.Join(savePath, defaultPublicKeyName)

	privateKey, err := generatePrivateKey(int(bitSize))
	if err != nil {
		return nil, err
	}

	publicKeyBytes, err := generatePublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, err
	}

	privateKeyBytes := encodePrivateKeyToPEM(privateKey)

	if err := writeKeyToFile(privateKeyBytes, pkPath); err != nil {
		return nil, err
	}

	if err := writeKeyToFile(publicKeyBytes, pubPath); err != nil {
		return nil, err
	}
	return privateKey, nil
}

func LoadPrivateKey(folderPath string) (*rsa.PrivateKey, error) {
	pkPath := path.Join(folderPath, defaultPrivateKeyName)
	f, err := os.ReadFile(pkPath)
	if err != nil {
		return nil, err
	}
	block, _ := pem.Decode(f)
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return key, nil
}

func generatePrivateKey(bitSize int) (*rsa.PrivateKey, error) {
	// Private Key generation
	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		return nil, err
	}

	// Validate Private Key
	err = privateKey.Validate()
	if err != nil {
		return nil, err
	}

	log.Println("Private Key generated")
	return privateKey, nil
}

// encodePrivateKeyToPEM encodes Private Key from RSA to PEM format
func encodePrivateKeyToPEM(privateKey *rsa.PrivateKey) []byte {
	// Get ASN.1 DER format
	privDER := x509.MarshalPKCS1PrivateKey(privateKey)

	// pem.Block
	privBlock := pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   privDER,
	}

	// Private key in PEM format
	privatePEM := pem.EncodeToMemory(&privBlock)

	return privatePEM
}

// generatePublicKey take a rsa.PublicKey and return bytes suitable for writing to .pub file
// returns in the format "ssh-rsa ..."
func generatePublicKey(privatekey *rsa.PublicKey) ([]byte, error) {
	publicRsaKey, err := ssh.NewPublicKey(privatekey)
	if err != nil {
		return nil, err
	}

	pubKeyBytes := ssh.MarshalAuthorizedKey(publicRsaKey)

	log.Println("Public key generated")
	return pubKeyBytes, nil
}

// writePemToFile writes keys to a file
func writeKeyToFile(keyBytes []byte, saveFileTo string) error {
	err := os.WriteFile(saveFileTo, keyBytes, 0600)
	if err != nil {
		return err
	}

	log.Printf("Key saved to: %s", saveFileTo)
	return nil
}