quorate/services/short/vault/client.go

package vault

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"
)

type Client struct {
	URL, Token string
	HTTP       *http.Client
}
type Record struct {
	Data    map[string]any
	Version int
}

func Login(ctx context.Context, url, role, tokenFile string) (*Client, error) {
	jwt, err := os.ReadFile(tokenFile)
	if err != nil {
		return nil, err
	}
	c := &Client{URL: strings.TrimRight(url, "/"), HTTP: &http.Client{Timeout: 10 * time.Second}}
	var response struct {
		Auth struct {
			Token string `json:"client_token"`
		} `json:"auth"`
	}
	err = c.call(ctx, http.MethodPost, "auth/kubernetes/login", map[string]string{"role": role, "jwt": strings.TrimSpace(string(jwt))}, &response)
	if err != nil {
		return nil, err
	}
	c.Token = response.Auth.Token
	return c, nil
}

func (c *Client) call(ctx context.Context, method, path string, input, output any) error {
	var body io.Reader
	if input != nil {
		b, err := json.Marshal(input)
		if err != nil {
			return err
		}
		body = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(ctx, method, c.URL+"/v1/"+path, body)
	if err != nil {
		return err
	}
	if input != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.Token != "" {
		req.Header.Set("X-Vault-Token", c.Token)
	}
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("Vault %s HTTP %d", path, resp.StatusCode)
	}
	if output != nil {
		return json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(output)
	}
	return nil
}

func (c *Client) Read(ctx context.Context, path string) (Record, error) {
	var response struct {
		Data struct {
			Values   map[string]any `json:"data"`
			Metadata struct {
				Version int `json:"version"`
			} `json:"metadata"`
		} `json:"data"`
	}
	err := c.call(ctx, http.MethodGet, "secret/data/"+path, nil, &response)
	return Record{response.Data.Values, response.Data.Metadata.Version}, err
}

func (c *Client) Write(ctx context.Context, path string, record Record) error {
	return c.call(ctx, http.MethodPost, "secret/data/"+path, map[string]any{"options": map[string]int{"cas": record.Version}, "data": record.Data}, nil)
}

func (c *Client) Seed(ctx context.Context, path string) error {
	r, err := c.Read(ctx, path)
	if err != nil {
		return err
	}
	value, _ := r.Data["private_key"].(string)
	var key *rsa.PrivateKey
	if value != "" && value != "PLACEHOLDER" {
		block, _ := pem.Decode([]byte(value))
		if block == nil {
			return fmt.Errorf("invalid API key encoding")
		}
		key, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return fmt.Errorf("invalid API private key")
		}
	} else {
		key, err = rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return err
		}
		r.Data = map[string]any{"private_key": string(pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)}))}
		if err = c.Write(ctx, path, r); err != nil {
			return err
		}
	}
	pub, err := x509.MarshalPKIXPublicKey(&key.PublicKey)
	if err != nil {
		return err
	}
	public := string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pub}))
	r, err = c.Read(ctx, "short/public")
	if err != nil {
		return err
	}
	if r.Data["public_key"] == public {
		return nil
	}
	r.Data = map[string]any{"public_key": public}
	return c.Write(ctx, "short/public", r)
}