alert_upload_encryption_example.go

package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io"
	"log"
	"os"
)

type Encryptor struct {
	pub   *rsa.PublicKey
	keyID string
}

type uploadAlertDataHeader struct {
	Version             int    `json:"version"`
	KeyID               string `json:"key_id"`
	EncryptedSessionKey string `json:"encrypted_session_key"`
	Iv                  string `json:"iv"`
	AuthTag             string `json:"auth_tag"`
}

func makeEncryptor(pubKeyPEM string, keyID string) (*Encryptor, error) {
	block, _ := pem.Decode([]byte(pubKeyPEM))
	if block == nil {
		return nil, fmt.Errorf("cannot decode public key: %s", pubKeyPEM)
	}

	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	switch pub := pub.(type) {
	case *rsa.PublicKey:
		return &Encryptor{pub, keyID}, nil
	default:
		return nil, fmt.Errorf("public key wasn't RSA key, was %T", pub)
	}
}

func (e *Encryptor) wrapSessionKey(sessionKey []byte) ([]byte, error) {
	encryptedKey, err := rsa.EncryptPKCS1v15(rand.Reader, e.pub, sessionKey)
	if err != nil {
		return nil, err
	}

	return encryptedKey, nil
}

func (e *Encryptor) generateSessionCipher() (cipher.AEAD, []byte, error) {
	sessionKeySize := 32
	sessionKey := make([]byte, sessionKeySize)
	_, err := rand.Read(sessionKey)
	if err != nil {
		return nil, nil, err
	}

	block, err := aes.NewCipher(sessionKey)
	if err != nil {
		log.Fatal(err.Error())
	}

	sessionCipher, err := cipher.NewGCM(block)
	if err != nil {
		return nil, nil, err
	}

	return sessionCipher, sessionKey, nil
}

func (e *Encryptor) Encrypt(plaintext []byte) ([]byte, error) {
	sessionCipher, sessionKey, err := e.generateSessionCipher()
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, sessionCipher.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}

	tagLen := 16 // this is a private constant called `gcmTagSize` found in crypto/cipher/gcm.go
	cipherTextWithTag := sessionCipher.Seal(nil, nonce, plaintext, nil)

	// NB: strip the auth tag off the end of the ciphertext. some libraries append it, some emit it separately. golang appends it.
	tagIdx := len(cipherTextWithTag) - tagLen
	authTag := cipherTextWithTag[tagIdx:]
	cipherText := cipherTextWithTag[:tagIdx]
	wrappedSessionKey, err := e.wrapSessionKey(sessionKey)
	if err != nil {
		return nil, err
	}

	header := uploadAlertDataHeader{
		Version:             1,
		KeyID:               e.keyID,
		EncryptedSessionKey: base64.StdEncoding.EncodeToString(wrappedSessionKey),
		Iv:                  base64.StdEncoding.EncodeToString(nonce),
		AuthTag:             base64.StdEncoding.EncodeToString(authTag),
	}
	headerJSON, err := json.Marshal(header)
	if err != nil {
		return nil, err
	}

	content := append(headerJSON, append([]byte("\x00"), cipherText...)...)
	return content, nil
}

func main() {
	// These are just hard-coded examples. Use the ones provided to you!
	pubKeyPEM := "-----BEGIN PUBLIC KEY-----\nMIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEA1EyQCSwyiNm8AZd6RDY6\njR249uYfNw94zkNla59xvWSrewxpFJZIAcTYRKwu0PjqmDPGYgqCcqM1Mfy9g9Ih\n27A1HNjn4wktvQDZSBtolLEpiyQgJLxBkcy9rlg+Dol8SadZxH+au5lIB9F6ghX8\ncqRFIjAjtob/L4aJA4OhC+GnjuDOEBnyVpHKIj0BlRJ7O0wyk41FMJvsaOnZsRrF\nrnp8C7zvG10In/gv+Wbow6md8fAn1+XxGrJGkl/ug+Z1pfLaoLCkOvz3RCRn92iX\niKgQfyme3OuKdSZjN6LafYnroo+HhlyaL9lYfYg1eJftLTwHugH4vE7ACw18AACs\nLcIQ6f0gWM2boWP9u++/KUEpo0g6WKOMpPFUcndvpTMaA9/FIJ5QJNEOJ8j99p//\nzGZRCLw6P/D86OVWdm0JCoryFQgmGNw7TwIMR9SaSdxKanmhGr80obxDk4Hd+5Vd\n1e9Lhj0XXDYcSXOpZC3a8LSSTElDijwpEvC08r/94GY7UiHFHwajAYoxZnf9PbWj\n493gR7e1Nzf276wROCSRvSBhUEnT76QrOOG04CqTN4hocuQfIGfAhyHPCeL4aCTs\npTz+rDWmeceLTeKPe2d9IWkJaf9mrO3da6h8auciJ7QB6YyjC9m//HPv/Gu33+8X\nNCjedLEmFA1HHs71ojFSXksCAwEAAQ==\n-----END PUBLIC KEY-----\n"
	keyID := "uLpzPZqTQN3nR248fgNJs3CU"

	encryptor, err := makeEncryptor(pubKeyPEM, keyID)
	if err != nil {
		log.Fatal(err)
	}

	// Your Alert JSON goes here.
	encrypted, err := encryptor.Encrypt([]byte("interesting plaintext JSON"))
	if err != nil {
		log.Fatal(err)
	}

	// for the sake of demonstration, write the encrypted envelope out to a tempfile
	f, err := os.CreateTemp("", "")
	if err != nil {
		log.Fatal(err)
	}
	if _, err := f.Write(encrypted); err != nil {
		log.Fatal(err)
	}
	if err := f.Close(); err != nil {
		log.Fatal(err)
	}
	fmt.Printf("output: %s", f.Name())
}