github.com/bshelton229/agent@v3.5.4+incompatible/agent/gs_downloader.go (about)

     1  package agent
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"strings"
     8  
     9  	"golang.org/x/oauth2/google"
    10  	storage "google.golang.org/api/storage/v1"
    11  )
    12  
    13  type GSDownloader struct {
    14  	// The name of the bucket
    15  	Bucket string
    16  
    17  	// The root directory of the download
    18  	Destination string
    19  
    20  	// The relative path that should be preserved in the download folder,
    21  	// also it's location in the bucket
    22  	Path string
    23  
    24  	// How many times should it retry the download before giving up
    25  	Retries int
    26  
    27  	// If failed responses should be dumped to the log
    28  	DebugHTTP bool
    29  }
    30  
    31  func (d GSDownloader) Start() error {
    32  	client, err := google.DefaultClient(context.Background(), storage.DevstorageReadOnlyScope)
    33  	if err != nil {
    34  		return errors.New(fmt.Sprintf("Error creating Google Cloud Storage client: %v", err))
    35  	}
    36  
    37  	url := "https://www.googleapis.com/storage/v1/b/" + d.BucketName() + "/o/" + escape(d.BucketFileLocation()) + "?alt=media"
    38  
    39  	// We can now cheat and pass the URL onto our regular downloader
    40  	return Download{
    41  		Client:      *client,
    42  		URL:         url,
    43  		Path:        d.Path,
    44  		Destination: d.Destination,
    45  		Retries:     d.Retries,
    46  		DebugHTTP:   d.DebugHTTP,
    47  	}.Start()
    48  }
    49  
    50  func (d GSDownloader) BucketFileLocation() string {
    51  	if d.BucketPath() != "" {
    52  		return strings.TrimSuffix(d.BucketPath(), "/") + "/" + strings.TrimPrefix(d.Path, "/")
    53  	} else {
    54  		return d.Path
    55  	}
    56  }
    57  
    58  func (d GSDownloader) BucketPath() string {
    59  	return strings.Join(d.destinationParts()[1:len(d.destinationParts())], "/")
    60  }
    61  
    62  func (d GSDownloader) BucketName() string {
    63  	return d.destinationParts()[0]
    64  }
    65  
    66  func (d GSDownloader) destinationParts() []string {
    67  	trimmed := strings.TrimPrefix(d.Bucket, "gs://")
    68  
    69  	return strings.Split(trimmed, "/")
    70  }
    71  
    72  func escape(s string) string {
    73  	// See https://golang.org/src/net/url/url.go
    74  	hexCount := 0
    75  	for i := 0; i < len(s); i++ {
    76  		c := s[i]
    77  		if shouldEscape(c) {
    78  			hexCount++
    79  		}
    80  	}
    81  
    82  	if hexCount == 0 {
    83  		return s
    84  	}
    85  
    86  	t := make([]byte, len(s)+2*hexCount)
    87  	j := 0
    88  	for i := 0; i < len(s); i++ {
    89  		switch c := s[i]; {
    90  		case shouldEscape(c):
    91  			t[j] = '%'
    92  			t[j+1] = "0123456789ABCDEF"[c>>4]
    93  			t[j+2] = "0123456789ABCDEF"[c&15]
    94  			j += 3
    95  		default:
    96  			t[j] = s[i]
    97  			j++
    98  		}
    99  	}
   100  	return string(t)
   101  }
   102  
   103  func shouldEscape(c byte) bool {
   104  	// See https://cloud.google.com/storage/docs/json_api/#encoding
   105  	if 'A' <= c && c <= 'Z' || 'a' <= c && c <= 'z' || '0' <= c && c <= '9' {
   106  		return false
   107  	}
   108  	switch c {
   109  	case '-', '.', '_', '~', '!', '$', '&', '\'', '(', ')', '*', '+', ',', ';', '=', ':', '@':
   110  		return false
   111  	}
   112  	return true
   113  }