security: implement prompt injection defense and credential encryption

cmd/picoclaw/internal/helpers.go

@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"runtime"
 
+	"github.com/sipeed/picoclaw/pkg/auth"
 	"github.com/sipeed/picoclaw/pkg/config"
 )
 
@@ -24,7 +25,26 @@ func GetConfigPath() string {
 }
 
 func LoadConfig() (*config.Config, error) {
-	return config.LoadConfig(GetConfigPath())
+	cfg, err := config.LoadConfig(GetConfigPath())
+	if err != nil {
+		return nil, err
+	}
+
+	// Initialize secure store with config settings
+	if err := initSecureStore(cfg); err != nil {
+		return nil, fmt.Errorf("initializing secure store: %w", err)
+	}
+
+	return cfg, nil
+}
+
+// initSecureStore initializes the secure credential store based on config.
+func initSecureStore(cfg *config.Config) error {
+	return auth.InitSecureStore(auth.SecureStoreConfig{
+		Enabled:     cfg.Security.CredentialEncryption.Enabled,
+		UseKeychain: cfg.Security.CredentialEncryption.UseKeychain,
+		Algorithm:   cfg.Security.CredentialEncryption.Algorithm,
+	})
 }
 
 // FormatVersion returns the version string with optional git commit

go.mod

@@ -18,11 +18,15 @@ require (
 	github.com/spf13/cobra v1.10.2
 	github.com/stretchr/testify v1.11.1
 	github.com/tencent-connect/botgo v0.2.1
+	github.com/zalando/go-keyring v0.2.6
 	golang.org/x/oauth2 v0.35.0
 )
 
 require (
+	al.essio.dev/pkg/shellescape v1.5.1 // indirect
+	github.com/danieljoos/wincred v1.2.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/godbus/dbus/v5 v5.1.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/spf13/pflag v1.0.10 // indirect
@@ -51,7 +55,7 @@ require (
 	github.com/valyala/fasthttp v1.69.0 // indirect
 	github.com/valyala/fastjson v1.6.7 // indirect
 	golang.org/x/arch v0.24.0 // indirect
-	golang.org/x/crypto v0.48.0 // indirect
+	golang.org/x/crypto v0.48.0
 	golang.org/x/net v0.50.0 // indirect
 	golang.org/x/sync v0.19.0 // indirect
 	golang.org/x/sys v0.41.0 // indirect

go.sum

@@ -1,3 +1,5 @@
+al.essio.dev/pkg/shellescape v1.5.1 h1:86HrALUujYS/h+GtqoB26SBEdkWfmMI6FubjXlsXyho=
+al.essio.dev/pkg/shellescape v1.5.1/go.mod h1:6sIqp7X2P6mThCQ7twERpZTuigpr6KbZWtls1U8I890=
 cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k=
 github.com/adhocore/gronx v1.19.6 h1:5KNVcoR9ACgL9HhEqCm5QXsab/gI4QDIybTAWcXDKDc=
 github.com/adhocore/gronx v1.19.6/go.mod h1:7oUY1WAU8rEJWmAxXR2DN0JaO4gi9khSgKjiRypqteg=
@@ -27,6 +29,8 @@ github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI
 github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
+github.com/danieljoos/wincred v1.2.2 h1:774zMFJrqaeYCK2W57BgAem/MLi6mtSE47MB6BOJ0i0=
+github.com/danieljoos/wincred v1.2.2/go.mod h1:w7w4Utbrz8lqeMbDAK0lkNJUv5sAOkFi7nd/ogr0Uh8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -42,6 +46,8 @@ github.com/go-resty/resty/v2 v2.17.1/go.mod h1:kCKZ3wWmwJaNc7S29BRtUhJwy7iqmn+2m
 github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0/go.mod h1:fyg7847qk6SyHyPtNmDHnmrv/HOrqktSC+C9fM+CJOE=
 github.com/go-test/deep v1.1.1 h1:0r/53hagsehfO4bzD2Pgr/+RgHqhmf+k1Bpse2cTu1U=
 github.com/go-test/deep v1.1.1/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE=
+github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk=
+github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -63,6 +69,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8=
 github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
@@ -122,6 +130,7 @@ github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3A
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
@@ -158,6 +167,8 @@ github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3i
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
+github.com/zalando/go-keyring v0.2.6 h1:r7Yc3+H+Ux0+M72zacZoItR3UDxeWfKTcabvkI8ua9s=
+github.com/zalando/go-keyring v0.2.6/go.mod h1:2TCrxYrbUNYfNS/Kgy/LSrkSQzZ5UPVH85RwfczwvcI=
 go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y=
 go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU=
 go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=

pkg/auth/encryption.go

@@ -0,0 +1,279 @@
+package auth
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+
+	"golang.org/x/crypto/chacha20poly1305"
+)
+
+var (
+	ErrEncryptionFailed     = errors.New("encryption failed")
+	ErrDecryptionFailed     = errors.New("decryption failed")
+	ErrInvalidCiphertext    = errors.New("invalid ciphertext")
+	ErrKeyNotFound          = errors.New("encryption key not found")
+	ErrUnsupportedAlgorithm = errors.New("unsupported encryption algorithm")
+)
+
+// EncryptionAlgorithm defines the supported encryption algorithms.
+type EncryptionAlgorithm string
+
+const (
+	AlgorithmChaCha20Poly1305 EncryptionAlgorithm = "chacha20-poly1305"
+	AlgorithmAES256GCM        EncryptionAlgorithm = "aes-256-gcm"
+)
+
+// EncryptedData represents encrypted credential data with metadata.
+type EncryptedData struct {
+	Algorithm  string `json:"algorithm"`
+	Nonce      string `json:"nonce"`
+	Ciphertext string `json:"ciphertext"`
+}
+
+// Encryptor provides encryption/decryption functionality for credentials.
+type Encryptor struct {
+	algorithm EncryptionAlgorithm
+	key       []byte
+}
+
+// NewEncryptor creates a new encryptor with the specified algorithm.
+func NewEncryptor(algorithm string) (*Encryptor, error) {
+	alg := EncryptionAlgorithm(algorithm)
+	if alg != AlgorithmChaCha20Poly1305 && alg != AlgorithmAES256GCM {
+		return nil, fmt.Errorf("%w: %s", ErrUnsupportedAlgorithm, algorithm)
+	}
+
+	key, err := getOrCreateEncryptionKey(alg)
+	if err != nil {
+		return nil, fmt.Errorf("getting encryption key: %w", err)
+	}
+
+	return &Encryptor{
+		algorithm: alg,
+		key:       key,
+	}, nil
+}
+
+// Encrypt encrypts the given data and returns base64-encoded ciphertext.
+func (e *Encryptor) Encrypt(plaintext []byte) (*EncryptedData, error) {
+	switch e.algorithm {
+	case AlgorithmChaCha20Poly1305:
+		return e.encryptChaCha20Poly1305(plaintext)
+	case AlgorithmAES256GCM:
+		return e.encryptAES256GCM(plaintext)
+	default:
+		return nil, fmt.Errorf("%w: %s", ErrUnsupportedAlgorithm, e.algorithm)
+	}
+}
+
+// Decrypt decrypts the given encrypted data.
+func (e *Encryptor) Decrypt(data *EncryptedData) ([]byte, error) {
+	switch data.Algorithm {
+	case string(AlgorithmChaCha20Poly1305):
+		return e.decryptChaCha20Poly1305(data)
+	case string(AlgorithmAES256GCM):
+		return e.decryptAES256GCM(data)
+	default:
+		return nil, fmt.Errorf("%w: %s", ErrUnsupportedAlgorithm, data.Algorithm)
+	}
+}
+
+// EncryptCredential encrypts an AuthCredential struct.
+func (e *Encryptor) EncryptCredential(cred *AuthCredential) (*EncryptedData, error) {
+	data, err := json.Marshal(cred)
+	if err != nil {
+		return nil, fmt.Errorf("marshaling credential: %w", err)
+	}
+	return e.Encrypt(data)
+}
+
+// DecryptCredential decrypts encrypted data into an AuthCredential.
+func (e *Encryptor) DecryptCredential(encData *EncryptedData) (*AuthCredential, error) {
+	plaintext, err := e.Decrypt(encData)
+	if err != nil {
+		return nil, err
+	}
+
+	var cred AuthCredential
+	if err := json.Unmarshal(plaintext, &cred); err != nil {
+		return nil, fmt.Errorf("unmarshaling credential: %w", err)
+	}
+
+	return &cred, nil
+}
+
+func (e *Encryptor) encryptChaCha20Poly1305(plaintext []byte) (*EncryptedData, error) {
+	aead, err := chacha20poly1305.NewX(e.key)
+	if err != nil {
+		return nil, fmt.Errorf("%w: creating cipher: %v", ErrEncryptionFailed, err)
+	}
+
+	nonce := make([]byte, aead.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, fmt.Errorf("%w: generating nonce: %v", ErrEncryptionFailed, err)
+	}
+
+	ciphertext := aead.Seal(nil, nonce, plaintext, nil)
+
+	return &EncryptedData{
+		Algorithm:  string(AlgorithmChaCha20Poly1305),
+		Nonce:      base64.StdEncoding.EncodeToString(nonce),
+		Ciphertext: base64.StdEncoding.EncodeToString(ciphertext),
+	}, nil
+}
+
+func (e *Encryptor) decryptChaCha20Poly1305(data *EncryptedData) ([]byte, error) {
+	aead, err := chacha20poly1305.NewX(e.key)
+	if err != nil {
+		return nil, fmt.Errorf("%w: creating cipher: %v", ErrDecryptionFailed, err)
+	}
+
+	nonce, err := base64.StdEncoding.DecodeString(data.Nonce)
+	if err != nil {
+		return nil, fmt.Errorf("%w: decoding nonce: %v", ErrInvalidCiphertext, err)
+	}
+
+	ciphertext, err := base64.StdEncoding.DecodeString(data.Ciphertext)
+	if err != nil {
+		return nil, fmt.Errorf("%w: decoding ciphertext: %v", ErrInvalidCiphertext, err)
+	}
+
+	if len(nonce) != aead.NonceSize() {
+		return nil, fmt.Errorf("%w: invalid nonce size", ErrInvalidCiphertext)
+	}
+
+	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
+	if err != nil {
+		return nil, fmt.Errorf("%w: decrypting: %v", ErrDecryptionFailed, err)
+	}
+
+	return plaintext, nil
+}
+
+func (e *Encryptor) encryptAES256GCM(plaintext []byte) (*EncryptedData, error) {
+	block, err := aes.NewCipher(e.key)
+	if err != nil {
+		return nil, fmt.Errorf("%w: creating cipher: %v", ErrEncryptionFailed, err)
+	}
+
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, fmt.Errorf("%w: creating GCM: %v", ErrEncryptionFailed, err)
+	}
+
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, fmt.Errorf("%w: generating nonce: %v", ErrEncryptionFailed, err)
+	}
+
+	ciphertext := gcm.Seal(nil, nonce, plaintext, nil)
+
+	return &EncryptedData{
+		Algorithm:  string(AlgorithmAES256GCM),
+		Nonce:      base64.StdEncoding.EncodeToString(nonce),
+		Ciphertext: base64.StdEncoding.EncodeToString(ciphertext),
+	}, nil
+}
+
+func (e *Encryptor) decryptAES256GCM(data *EncryptedData) ([]byte, error) {
+	block, err := aes.NewCipher(e.key)
+	if err != nil {
+		return nil, fmt.Errorf("%w: creating cipher: %v", ErrDecryptionFailed, err)
+	}
+
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, fmt.Errorf("%w: creating GCM: %v", ErrDecryptionFailed, err)
+	}
+
+	nonce, err := base64.StdEncoding.DecodeString(data.Nonce)
+	if err != nil {
+		return nil, fmt.Errorf("%w: decoding nonce: %v", ErrInvalidCiphertext, err)
+	}
+
+	ciphertext, err := base64.StdEncoding.DecodeString(data.Ciphertext)
+	if err != nil {
+		return nil, fmt.Errorf("%w: decoding ciphertext: %v", ErrInvalidCiphertext, err)
+	}
+
+	if len(nonce) != gcm.NonceSize() {
+		return nil, fmt.Errorf("%w: invalid nonce size", ErrInvalidCiphertext)
+	}
+
+	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
+	if err != nil {
+		return nil, fmt.Errorf("%w: decrypting: %v", ErrDecryptionFailed, err)
+	}
+
+	return plaintext, nil
+}
+
+// getOrCreateEncryptionKey retrieves or creates an encryption key for file-based encryption.
+// This is used as a fallback when OS keychain is not available.
+func getOrCreateEncryptionKey(algorithm EncryptionAlgorithm) ([]byte, error) {
+	keySize := 32 // Both ChaCha20-Poly1305 and AES-256 use 32-byte keys
+
+	keyPath, err := encryptionKeyPath()
+	if err != nil {
+		return nil, err
+	}
+
+	// Try to read existing key
+	key, err := os.ReadFile(keyPath)
+	if err == nil && len(key) == keySize {
+		return key, nil
+	}
+
+	// Generate new key
+	key = make([]byte, keySize)
+	if _, err := io.ReadFull(rand.Reader, key); err != nil {
+		return nil, fmt.Errorf("generating key: %w", err)
+	}
+
+	// Ensure directory exists
+	dir := filepath.Dir(keyPath)
+	if err := os.MkdirAll(dir, 0o700); err != nil {
+		return nil, fmt.Errorf("creating key directory: %w", err)
+	}
+
+	// Write key with restricted permissions
+	if err := os.WriteFile(keyPath, key, 0o600); err != nil {
+		return nil, fmt.Errorf("writing key: %w", err)
+	}
+
+	return key, nil
+}
+
+func encryptionKeyPath() (string, error) {
+	home := os.Getenv("HOME")
+	if home == "" {
+		var err error
+		home, err = os.UserHomeDir()
+		if err != nil {
+			return "", fmt.Errorf("getting home dir: %w", err)
+		}
+	}
+	return filepath.Join(home, ".picoclaw", ".key"), nil
+}
+
+// DeleteEncryptionKey removes the encryption key file.
+// This should be called when all credentials are deleted.
+func DeleteEncryptionKey() error {
+	keyPath, err := encryptionKeyPath()
+	if err != nil {
+		return err
+	}
+
+	if err := os.Remove(keyPath); err != nil && !os.IsNotExist(err) {
+		return err
+	}
+	return nil
+}