Automatically download Astra secure connect bundle

astra/bundle.go

@@ -16,15 +16,24 @@ package astra
 
 import (
 	"archive/zip"
+	"bytes"
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net/http"
 	"runtime"
+	"time"
+
+	"github.com/datastax/astra-client-go/v2/astra"
 )
 
+const URL = "https://api.astra.datastax.com"
+
 type Bundle struct {
 	tlsConfig *tls.Config
 	host      string
@@ -85,6 +94,68 @@ func LoadBundleZipFromPath(path string) (*Bundle, error) {
 	return LoadBundleZip(&reader.Reader)
 }
 
+func LoadBundleZipFromURL(url, databaseID, token string, timeout time.Duration) (*Bundle, error) {
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(timeout))
+	defer cancel()
+
+	credsURL, err := generateSecureBundleURLWithResponse(url, databaseID, token, ctx)
+	if err != nil {
+		return nil, fmt.Errorf("error generating secure bundle zip URLs: %v", err)
+	}
+	resp, err := http.Get(credsURL.DownloadURL)
+
+	defer resp.Body.Close()
+
+	body, err := readAllWithTimeout(resp.Body, ctx)
+	if err != nil {
+		return nil, fmt.Errorf("error downloading bundle zip: %v", err)
+	}
+
+	reader, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
+	if err != nil {
+		return nil, fmt.Errorf("error creating zip reader for bundle zip: %v", err)
+	}
+
+	return LoadBundleZip(reader)
+}
+
+func readAllWithTimeout(r io.Reader, ctx context.Context) (bytes []byte, err error) {
+	ch := make(chan struct{})
+
+	go func() {
+		bytes, err = ioutil.ReadAll(r)
+		close(ch)
+	}()
+
+	select {
+	case <-ch:
+	case <-ctx.Done():
+		return nil, errors.New("timeout reading data")
+	}
+
+	return bytes, err
+}
+
+func generateSecureBundleURLWithResponse(url, databaseID, token string, ctx context.Context) (*astra.CredsURL, error) {
+	client, err := astra.NewClientWithResponses(url, func(c *astra.Client) error {
+		c.RequestEditors = append(c.RequestEditors, func(ctx context.Context, req *http.Request) error {
+			req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+			return nil
+		})
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	res, err := client.GenerateSecureBundleURLWithResponse(ctx, astra.DatabaseIdParam(databaseID))
+
+	if res.StatusCode() != http.StatusOK {
+		return nil, fmt.Errorf("unable to generate bundle urls, failed with status code %d", res.StatusCode())
+	}
+
+	return res.JSON200, nil
+}
+
 func (b *Bundle) Host() string {
 	return b.host
 }

astra/bundle_test.go

@@ -24,15 +24,16 @@ import (
 	"encoding/asn1"
 	"encoding/json"
 	"encoding/pem"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"io"
 	"io/ioutil"
 	"math/big"
 	"net"
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestLoadBundleZip(t *testing.T) {

go.mod

@@ -4,7 +4,7 @@ go 1.16
 
 require (
 	github.com/alecthomas/kong v0.2.17
-	github.com/antlr/antlr4/runtime/Go/antlr v0.0.0-20210521184019-c5ad59b459ec
+	github.com/datastax/astra-client-go/v2 v2.2.9 // indirect
 	github.com/datastax/go-cassandra-native-protocol v0.0.0-20211124104234-f6aea54fa801
 	github.com/hashicorp/golang-lru v0.5.4
 	github.com/stretchr/testify v1.7.0