RSA Private/Public Key Generator, Encrypt and Decrypt in Golang

keyield.go

package main

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
)

type Keyield struct {
	Private []byte
	Public  []byte
}

// generator private/public
func (k *Keyield) Create(n int) error {
	pri, err := rsa.GenerateKey(rand.Reader, n)
	if err != nil {
		return err
	}
	s := x509.MarshalPKCS1PrivateKey(pri)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: s,
	}
	k.Private = pem.EncodeToMemory(block)
	pub := &pri.PublicKey
	d, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		return err
	}
	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: d,
	}
	k.Public = pem.EncodeToMemory(block)
	return nil
}

// read from file(s)
func (k *Keyield) Read(name string) {
	flpath := name
	isPri := strings.HasSuffix(name, ".pri")
	isPub := strings.HasSuffix(name, ".pub")
	if !isPri && !isPub {
		flpath = fmt.Sprintf("%s_RSA.pri", name)
	}
	if isPri || !isPub {
		if _, err := os.Stat(flpath); !os.IsNotExist(err) {
			pri, err := ioutil.ReadFile(flpath)
			if err == nil {
				k.Private = pri
			}
		}
		if isPri {
			return
		}
	}

	if !isPri && !isPub {
		flpath = fmt.Sprintf("%s.pub", name)
	}
	if _, err := os.Stat(flpath); !os.IsNotExist(err) {
		pub, err := ioutil.ReadFile(flpath)
		if err == nil {
			k.Public = pub
		}
	}
}

// write to file(s)
func (k *Keyield) Write(name string) error {
	if len(k.Private) == 0 {
		err := k.Create(1024)
		if err != nil {
			return err
		}
	}

	flpath := name
	isPri := strings.HasSuffix(name, ".pri")
	isPub := strings.HasSuffix(name, ".pub")
	if !isPub {
		if !isPri {
			flpath = fmt.Sprintf("%s_RSA.pri", name)
		}
	}
	if isPri || !isPub {
		err := ioutil.WriteFile(flpath, k.Private, 0644)
		if isPri {
			return err
		}
	}

	if !isPub {
		flpath = fmt.Sprintf("%s.pub", name)
	}
	return ioutil.WriteFile(flpath, k.Public, 0644)
}

// encrypt data
func (k *Keyield) Encrypt(data []byte) ([]byte, error) {
	var result []byte
	block, _ := pem.Decode(k.Public)
	if block == nil {
		return result, errors.New("public key error")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return result, err
	}
	return rsa.EncryptPKCS1v15(rand.Reader, pub.(*rsa.PublicKey), data)
}

// decrypt data
func (k *Keyield) Decrypt(data []byte) ([]byte, error) {
	var result []byte
	block, _ := pem.Decode(k.Private)
	if block == nil {
		return result, errors.New("private key error")
	}
	pri, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return result, err
	}
	return rsa.DecryptPKCS1v15(rand.Reader, pri, data)
}

// verify signature
func (k *Keyield) Verify(data []byte) ([]byte, error) {
	var signature []byte
	h := sha256.New()
	h.Write(data)
	block, _ := pem.Decode(k.Private)
	if block == nil {
		return signature, errors.New("private key error")
	}
	pri, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return signature, err
	}

	signature, err = rsa.SignPKCS1v15(rand.Reader, pri, crypto.SHA256, h.Sum(nil))
	if err != nil {
		return signature, err
	}

	block, _ = pem.Decode(k.Public)
	if block == nil {
		return signature, err
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return signature, err
	}

	hashed := sha256.Sum256(data)
	err = rsa.VerifyPKCS1v15(pub.(*rsa.PublicKey), crypto.SHA256, hashed[:], signature)
	return signature, err
}

// -------------------------
// usage demo
// go run keyield.go
// go run keyield.go read
// go run keyield.go my-new-secret
// go run keyield.go read
func main() {
	if len(os.Args) == 2 {

		if os.Args[1] == "read" {
			var read = func(flpath string) {
				fmt.Println("--- read and decrypt demo -----")
				body, err := ioutil.ReadFile(flpath)
				if err != nil {
					log.Fatal(err)
				}
				k := Keyield{}
				k.Read("mindon")
				message, err := k.Decrypt(body)
				if err != nil {
					log.Fatal(err)
				}
				_, err = k.Verify(message)
				if err != nil {
					log.Fatal(err)
				}
				fmt.Println(string(message))
			}
			read("hello.secret")
			return
		}

		var write = func(data []byte) {
			k := Keyield{}
			k.Read("mindon.pub")
			body, err := k.Encrypt(data)
			if err != nil {
				log.Fatal(err)
			}
			err = ioutil.WriteFile("hello.secret", body, 0644)
			if err != nil {
				log.Fatal(err)
			}
		}
		write([]byte(os.Args[1]))
		return
	}

	k := Keyield{}
	err := k.Write("mindon")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(string(k.Private))
	fmt.Println(string(k.Public))

	data := []byte(`Hello 世界！
Mindon, 2022-02-22 11:22`)
	cliper, err := k.Encrypt(data)
	if err != nil {
		log.Fatal(err)
	}

	err = ioutil.WriteFile("hello.secret", cliper, 0644)
	if err != nil {
		fmt.Println(err)
	}
	fmt.Println(hex.EncodeToString(cliper))
	text, err := k.Decrypt(cliper)
	if err != nil {
		log.Println("decrypt failed")
		log.Fatal(err)
	}

	_, err = k.Verify(text)
	if err != nil {
		log.Println("signature verify failed")
		log.Fatal(err)
	}
	fmt.Println(string(text))
}