Skip to content

Commit

Permalink
Add AES encryption and decryption commands for inventory management
Browse files Browse the repository at this point in the history
  • Loading branch information
struckchure committed Dec 30, 2024
1 parent 238fb2c commit 2ae1bdf
Show file tree
Hide file tree
Showing 2 changed files with 145 additions and 1 deletion.
64 changes: 64 additions & 0 deletions cmd/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,60 @@ var agentRunWorkflowCmd = &cobra.Command{
},
}

var inventoryCmd = &cobra.Command{
Use: "inventory",
Short: "Manage inventory with AES encryption",
}

var encryptInventoryCmd = &cobra.Command{
Use: "encrypt",
Run: func(cmd *cobra.Command, args []string) {
inventoryFile, _ := cmd.Flags().GetString("inventory")
encryptionKey, _ := cmd.Flags().GetString("encryption-key")

encryptedInventory, err := storm.NewInventory().Encrypt(inventoryFile, encryptionKey)
if err != nil {
fmt.Println(err)

os.Exit(1)
}

fmt.Println(*encryptedInventory)
},
}

var decryptInventoryCmd = &cobra.Command{
Use: "decrypt",
Run: func(cmd *cobra.Command, args []string) {
encryptedInventory, _ := cmd.Flags().GetString("encrypted-inventory")
encryptionKey, _ := cmd.Flags().GetString("encryption-key")
format, _ := cmd.Flags().GetString("format")

var byteEncryptedInventory []byte
var err error

if format == "file" {
byteEncryptedInventory, err = os.ReadFile(encryptedInventory)
if err != nil {
fmt.Println(err)

os.Exit(1)
}
} else {
byteEncryptedInventory = []byte(encryptedInventory)
}

decryptedInventory, err := storm.NewInventory().Decrypt(string(byteEncryptedInventory), encryptionKey)
if err != nil {
fmt.Println(err)

os.Exit(1)
}

fmt.Println(*decryptedInventory)
},
}

var agentInstallCmd = &cobra.Command{
Use: "install",
Run: func(cmd *cobra.Command, args []string) {
Expand Down Expand Up @@ -128,6 +182,16 @@ func main() {
agentRunWorkflowCmd.Flags().IntP("format", "f", 1, "available options are; 1 => plain, 2 => struct, 3 => json")
agentCmd.AddCommand(agentRunWorkflowCmd)

inventoryCmd.AddCommand(encryptInventoryCmd)
encryptInventoryCmd.Flags().StringP("inventory", "i", "./inventory.yaml", "formatio storm inventory")

inventoryCmd.AddCommand(decryptInventoryCmd)
decryptInventoryCmd.Flags().StringP("encrypted-inventory", "e", "./inventory.yaml.enc", "encrypted inventory file")
decryptInventoryCmd.Flags().StringP("format", "f", "file", "available options are; plain, file")

inventoryCmd.PersistentFlags().StringP("encryption-key", "k", "", "encryption key")
rootCmd.AddCommand(inventoryCmd)

runWorkflowCmd.Flags().BoolP("trash-workflow", "t", true, "remove workflow file if the workflow is complete")
runWorkflowCmd.Flags().StringP("directory", "d", ".", "directory to run the workflow from")
runWorkflowCmd.Flags().IntP("format", "f", 1, "available options are; 1 => plain, 2 => struct, 3 => json")
Expand Down
82 changes: 81 additions & 1 deletion inventory.go
Original file line number Diff line number Diff line change
@@ -1,14 +1,20 @@
package storm

import (
"crypto/aes"
"crypto/cipher"
"crypto/rand"
"encoding/base64"
"errors"
"io"
"os"

"gopkg.in/yaml.v3"
)

type Inventory struct{}

func (c *Inventory) Load(file string) (*InventoryConfig, error) {
func (i *Inventory) Load(file string) (*InventoryConfig, error) {
fileContent, err := os.ReadFile(file)
if err != nil {
return nil, err
Expand All @@ -24,6 +30,80 @@ func (c *Inventory) Load(file string) (*InventoryConfig, error) {
return config, nil
}

func (i *Inventory) aes(encryptionKey string) (cipher.AEAD, error) {
key := []byte(encryptionKey)
if len(key) != 16 && len(key) != 24 && len(key) != 32 {
return nil, errors.New("encryption key must be 16, 24, or 32 bytes long")
}

block, err := aes.NewCipher(key)
if err != nil {
return nil, err
}

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

return gcm, nil
}

func (i *Inventory) Encrypt(file string, encryptionKey string) (*string, error) {
fileContent, err := os.ReadFile(file)
if err != nil {
return nil, err
}

gcm, err := i.aes(encryptionKey)
if err != nil {
return nil, err
}

// Create a nonce of the correct size
nonce := make([]byte, gcm.NonceSize())
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
return nil, err
}

// Encrypt and prepend nonce
ciphertext := gcm.Seal(nonce, nonce, fileContent, nil)

// Encode the ciphertext to Base64 for safe storage
encodedCiphertext := base64.StdEncoding.EncodeToString(ciphertext)
return &encodedCiphertext, nil
}

func (i *Inventory) Decrypt(ciphertext string, decryptionKey string) (*string, error) {
gcm, err := i.aes(decryptionKey)
if err != nil {
return nil, err
}

// Decode the Base64 ciphertext
ciphertextBytes, err := base64.StdEncoding.DecodeString(ciphertext)
if err != nil {
return nil, err
}

nonceSize := gcm.NonceSize()
if len(ciphertextBytes) < nonceSize {
return nil, errors.New("ciphertext too short")
}

// Separate nonce and actual ciphertext
nonce, encryptedMessage := ciphertextBytes[:nonceSize], ciphertextBytes[nonceSize:]

// Decrypt the ciphertext
plaintext, err := gcm.Open(nil, nonce, encryptedMessage, nil)
if err != nil {
return nil, err
}

decryptedFileContent := string(plaintext)
return &decryptedFileContent, nil
}

func NewInventory() *Inventory {
return &Inventory{}
}

0 comments on commit 2ae1bdf

Please sign in to comment.