package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"log"
	"os"
)

func generateOrRetrieveRSAKeys(name string, size int) (io.Reader, []byte, error) {
	var pemprvkey bytes.Buffer
	var prvKey *rsa.PrivateKey
	fd, err := os.Open(name + ".pem")
	if os.IsNotExist(err) {
		log.Println("Private RSA key not found, generating a new one")

		fd, err = os.Create(name + ".pem")
		if err != nil {
			return nil, nil, fmt.Errorf("unable to open the file for writing: %w", err)
		}
		defer fd.Close()

		prvKey, err = rsa.GenerateKey(rand.Reader, size)
		if err != nil {
			return nil, nil, err
		}
		pkcs1prvkey := x509.MarshalPKCS1PrivateKey(prvKey)

		var tmpprvkey bytes.Buffer
		w := io.MultiWriter(&tmpprvkey, &pemprvkey)

		err = pem.Encode(w, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: pkcs1prvkey})
		if err != nil {
			return nil, nil, err
		}

		// Save it!
		_, err = tmpprvkey.WriteTo(fd)
		if err != nil {
			return nil, nil, err
		}
	} else if err != nil {
		return nil, nil, err
	} else {
		defer fd.Close()

		var tmpprvkey bytes.Buffer
		io.Copy(io.MultiWriter(&tmpprvkey, &pemprvkey), fd)

		decodedpem, _ := pem.Decode(tmpprvkey.Bytes())
		if decodedpem == nil {
			return nil, nil, fmt.Errorf("unable to decode PEM data")
		} else if decodedpem.Type != "RSA PRIVATE KEY" {
			return nil, nil, fmt.Errorf("not the expected PEM type: %s intead of RSA PRIVATE KEY", decodedpem.Type)
		}

		prvKey, err = x509.ParsePKCS1PrivateKey(decodedpem.Bytes)
	}

	pubkey := prvKey.Public()
	pkixpub, err := x509.MarshalPKIXPublicKey(pubkey.(*rsa.PublicKey))
	if err != nil {
		return nil, nil, err
	}

	pempubkey := pem.EncodeToMemory(
		&pem.Block{
			Type:  "PUBLIC KEY",
			Bytes: pkixpub,
		},
	)

	return &pemprvkey, pempubkey, nil
}