github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/rsa.go (about)

     1  // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     2  
     3  package runner
     4  
     5  import (
     6  	"crypto/rand"
     7  	"crypto/rsa"
     8  	"crypto/x509"
     9  	"encoding/pem"
    10  	"io/ioutil"
    11  	"os"
    12  
    13  	"github.com/go-stack/stack"
    14  	"github.com/jjeffery/kv"
    15  )
    16  
    17  // This file contains a number of abberviated functions for handling and using SSH
    18  // style RSA key pairs.  It enforces that the max 4096 bits will be used and a
    19  // passphrase also will be used.
    20  //
    21  // This code is used mainly in the tests.
    22  
    23  // encryptPrivateKeyToPEM takes a private key and generates a PEM block for it
    24  // that has been encrypted using the supplied password, pwd.
    25  func encryptPrivateKeyToPEM(privateKey *rsa.PrivateKey, pwd string) (block *pem.Block, err kv.Error) {
    26  	block = &pem.Block{
    27  		Type:  "RSA PRIVATE KEY",
    28  		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
    29  	}
    30  
    31  	if len(pwd) == 0 {
    32  		return nil, kv.NewError("passphrase must be specified").With("stack", stack.Trace().TrimRuntime())
    33  	}
    34  
    35  	encBlock, errGo := x509.EncryptPEMBlock(rand.Reader, block.Type, block.Bytes, []byte(pwd), x509.PEMCipherAES256)
    36  	if errGo != nil {
    37  		return nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
    38  	}
    39  	block = encBlock
    40  
    41  	return block, nil
    42  }
    43  
    44  // generatePrivateKey generates a rsa.PrivateKey and SSH style .pub file content
    45  func generatePrivateKey(pwd string) (privateKey *rsa.PrivateKey, privatePEM []byte, err kv.Error) {
    46  	if len(pwd) == 0 {
    47  		return nil, nil, kv.NewError("passphrase must be specified").With("stack", stack.Trace().TrimRuntime())
    48  	}
    49  
    50  	privateKey, errGo := rsa.GenerateKey(rand.Reader, 4096)
    51  	if errGo != nil {
    52  		return nil, nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
    53  	}
    54  
    55  	if errGo = privateKey.Validate(); errGo != nil {
    56  		return nil, nil, kv.Wrap(errGo).With("stack", stack.Trace().TrimRuntime())
    57  	}
    58  
    59  	block, err := encryptPrivateKeyToPEM(privateKey, pwd)
    60  	if err != nil {
    61  		return nil, nil, err
    62  	}
    63  
    64  	return privateKey, pem.EncodeToMemory(block), nil
    65  }
    66  
    67  // extractPublicPEM transforms a rsa.PublicKey to PEM style .pub file content
    68  func extractPublicPEM(privateKey *rsa.PrivateKey) (publicPEM []byte, err kv.Error) {
    69  	return pem.EncodeToMemory(
    70  		&pem.Block{
    71  			Type:  "RSA PUBLIC KEY",
    72  			Bytes: x509.MarshalPKCS1PublicKey(&privateKey.PublicKey),
    73  		}), nil
    74  }
    75  
    76  // generateKeyPair produces a 4096 bit PEM formatted public and password protected RSA key pair
    77  func GenerateKeyPair(pwd string) (privatePEM []byte, publicPEM []byte, err kv.Error) {
    78  	privateKey, privatePEM, err := generatePrivateKey(pwd)
    79  	if err != nil {
    80  		return nil, nil, err
    81  	}
    82  	publicPEM, err = extractPublicPEM(privateKey)
    83  	if err != nil {
    84  		return nil, nil, err
    85  	}
    86  	return privatePEM, publicPEM, nil
    87  }
    88  
    89  func WriteKeyToFile(keyBytes []byte, outputFN string) (err kv.Error) {
    90  	if len(keyBytes) == 0 {
    91  		return kv.NewError("empty key").With("output-file", outputFN).With("stack", stack.Trace().TrimRuntime())
    92  	}
    93  	if _, err := os.Stat(outputFN); err == nil || os.IsExist(err) {
    94  		return kv.NewError("file exists already").With("output-file", outputFN).With("stack", stack.Trace().TrimRuntime())
    95  	}
    96  
    97  	if errGo := ioutil.WriteFile(outputFN, keyBytes, 0600); errGo != nil {
    98  		return kv.Wrap(errGo).With("output-file", outputFN).With("stack", stack.Trace().TrimRuntime())
    99  	}
   100  
   101  	return nil
   102  }