agent-tools/internal/agent/vaultkv.go

package agent

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
)

const (
	// DefaultKVMount is the KV-v2 mount holding agent-facing secrets.
	DefaultKVMount = "kv"
)

var (
	// ErrVaultDenied marks a 403 from Vault (the AppRole policy lacks the capability).
	ErrVaultDenied = errors.New("permission denied")
	// ErrVaultNotFound marks a 404 from Vault (mount or secret does not exist).
	ErrVaultNotFound = errors.New("not found")
)

// VaultClient issues authenticated requests against Vault/OpenBao using a token
// obtained from the agent AppRole.
type VaultClient struct {
	Addr  string
	Token string
	HTTP  *http.Client
}

// NewVaultClient performs the AppRole login (role_id only, no secret_id) and
// returns a client bound to the resulting client_token.
func NewVaultClient(addr, roleID string) (*VaultClient, error) {
	token, err := approleLogin(addr, roleID)
	if err != nil {
		return nil, err
	}
	return &VaultClient{Addr: addr, Token: token, HTTP: httpClient}, nil
}

func (c *VaultClient) do(method, path string, body any, out any) error {
	var reader io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return err
		}
		reader = bytes.NewReader(b)
	}
	url := strings.TrimRight(c.Addr, "/") + path
	req, err := http.NewRequest(method, url, reader)
	if err != nil {
		return err
	}
	req.Header.Set("X-Vault-Token", c.Token)
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}

	resp, err := c.HTTP.Do(req)
	if err != nil {
		return fmt.Errorf("vault %s %s: %w", method, path, err)
	}
	defer func() { _ = resp.Body.Close() }()
	data, _ := io.ReadAll(resp.Body)

	switch {
	case resp.StatusCode == http.StatusForbidden:
		return fmt.Errorf("vault %s %s: %w", method, path, ErrVaultDenied)
	case resp.StatusCode == http.StatusNotFound:
		return fmt.Errorf("vault %s %s: %w", method, path, ErrVaultNotFound)
	case resp.StatusCode < 200 || resp.StatusCode >= 300:
		return fmt.Errorf("vault %s %s: HTTP %d: %s", method, path, resp.StatusCode, strings.TrimSpace(string(data)))
	}

	if out != nil && len(data) > 0 {
		// Response bodies here carry secret material, so decode failures never
		// echo the body.
		if err := json.Unmarshal(data, out); err != nil {
			return fmt.Errorf("vault %s %s: decoding response: %w", method, path, err)
		}
	}
	return nil
}

// kvDataPath builds the KV-v2 data path for a mount and secret path.
func kvDataPath(mount, path string) string {
	return "/v1/" + strings.Trim(mount, "/") + "/data/" + strings.Trim(path, "/")
}

// ReadKV returns the data map of a KV-v2 secret.
func (c *VaultClient) ReadKV(mount, path string) (map[string]any, error) {
	var out struct {
		Data struct {
			Data map[string]any `json:"data"`
		} `json:"data"`
	}
	if err := c.do(http.MethodGet, kvDataPath(mount, path), nil, &out); err != nil {
		return nil, err
	}
	if out.Data.Data == nil {
		return nil, fmt.Errorf("vault read %s/%s: secret has no data", mount, path)
	}
	return out.Data.Data, nil
}

// ReadKVOptional is ReadKV but treats a missing secret as empty: a 404 or a
// deleted version (data: null) yields an empty map, not an error. Callers that
// read-modify-write a path that may not exist yet use this.
func (c *VaultClient) ReadKVOptional(mount, path string) (map[string]any, error) {
	var out struct {
		Data struct {
			Data map[string]any `json:"data"`
		} `json:"data"`
	}
	if err := c.do(http.MethodGet, kvDataPath(mount, path), nil, &out); err != nil {
		if errors.Is(err, ErrVaultNotFound) {
			return map[string]any{}, nil
		}
		return nil, err
	}
	if out.Data.Data == nil {
		return map[string]any{}, nil
	}
	return out.Data.Data, nil
}

// WriteKV writes a KV-v2 secret and returns the version it created.
func (c *VaultClient) WriteKV(mount, path string, data map[string]string) (int, error) {
	fields := make(map[string]any, len(data))
	for k, v := range data {
		fields[k] = v
	}
	return c.WriteKVAny(mount, path, fields)
}

// WriteKVAny writes a KV-v2 secret whose fields are not all strings (so a
// read-modify-write can put back values it did not author) and returns the
// version it created.
func (c *VaultClient) WriteKVAny(mount, path string, data map[string]any) (int, error) {
	var out struct {
		Data struct {
			Version int `json:"version"`
		} `json:"data"`
	}
	body := map[string]any{"data": data}
	if err := c.do(http.MethodPost, kvDataPath(mount, path), body, &out); err != nil {
		return 0, err
	}
	if out.Data.Version == 0 {
		return 0, fmt.Errorf("vault write %s/%s: no version in response", mount, path)
	}
	return out.Data.Version, nil
}

// StringField returns the first non-empty string value among the given keys.
func StringField(data map[string]any, keys ...string) string {
	for _, k := range keys {
		if s, ok := data[k].(string); ok && s != "" {
			return s
		}
	}
	return ""
}