Add AES encryption and decryption commands for inventory management

Author: struckchure

cmd/main.go

@@ -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) {
@@ -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")

inventory.go

@@ -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
@@ -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{}
 }