-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
194 lines (162 loc) · 4.7 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
package main
import (
"context"
"fmt"
"io/ioutil"
"log"
"net/url"
"os"
"strings"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/awserr"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/aws/aws-sdk-go/service/s3/s3manager"
"github.com/urfave/cli"
)
const (
appName = "ccat"
appUsage = "cloud cat\n\n A simple CLI that streams objects from S3 to STDOUT"
appVersion = "0.1.0"
appUsageText = "ccat s3://your-bucket/your-key https://s3-us-west-2.amazonaws.com/your-bucket/your-other-key"
prefixS3 = "s3://"
prefixHTTPS = "https://"
)
var (
// Stdout is a drop in replacement for os.Stdout when we have to use the io.WriterAt interface
Stdout = stdout{}
)
func main() {
app := cli.NewApp()
app.Name = appName
app.Usage = appUsage
app.UsageText = appUsageText
app.Version = appVersion
app.HideHelp = true
app.HideVersion = true
// TODO:
// * add bytes range
// * handle "*" prefix and stream multiple objects sequentially
// * add a bytes flag for max number of bytes to scan
// * add a max number of objects flag. bytes or max which ever comes first
app.Flags = []cli.Flag{
cli.BoolFlag{
Name: "help, h",
Usage: "show this help message",
},
cli.BoolFlag{
Name: "verbose",
Usage: "log verbosely to stderr",
},
}
app.Action = do
err := app.Run(os.Args)
if err != nil {
fmt.Fprintf(os.Stderr, "there was an error trying to stream objects: %v\n", err)
os.Exit(1)
}
}
func do(c *cli.Context) error {
if c.Bool("help") || len(c.Args()) == 0 {
cli.ShowAppHelp(c)
return nil
}
if c.Bool("verbose") {
// Logging to stderr since we're streaming the object to STDOUT
log.SetOutput(os.Stderr)
} else {
log.SetOutput(ioutil.Discard)
}
err := validateCredentials()
if err != nil {
return fmt.Errorf("could not find any valid credentials")
}
ctx := context.Background()
return streamObjectsFromS3(ctx, c.Args())
}
func streamObjectsFromS3(ctx context.Context, objects []string) error {
for _, obj := range objects {
bucket, key, err := parseS3ObjectString(obj)
if err != nil {
return fmt.Errorf("could not parse %v into bucket and key: %v", obj, err)
}
log.Println("bucket:", bucket)
log.Println("key:", key)
err = streamObjectFromS3(ctx, bucket, key)
if err != nil {
return err
}
}
return nil
}
func streamObjectFromS3(ctx context.Context, bucket string, key string) error {
// figure out which region the bucket is in
log.Println("getting bucket region")
sess := session.Must(session.NewSession())
// TODO: consider caching the region for a bucket in memory
// TODO: the HTTP URLS have the region built in and we could pass the hint down here
region, err := s3manager.GetBucketRegion(ctx, sess, bucket, "us-east-1")
if err != nil {
if aerr, ok := err.(awserr.Error); ok && aerr.Code() == "NotFound" {
fmt.Fprintf(os.Stderr, "unable to find bucket %s's region not found\n", bucket)
}
return err
}
// TODO: consider caching sessions
log.Println("creating session for", region)
streamSess := session.Must(
session.NewSession(
&aws.Config{Region: aws.String(region)},
),
)
svc := s3manager.NewDownloader(streamSess, func(d *s3manager.Downloader) {
// Force the downloader to stream the object sequentially
d.Concurrency = 1
})
log.Printf("streaming %s/%s\n", bucket, key)
return stream(bucket, key, svc)
}
func stream(bucket string, key string, svc *s3manager.Downloader) error {
input := &s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
}
ctx := context.Background()
_, err := svc.DownloadWithContext(ctx, Stdout, input)
return err
}
// parseS3ObjectString can parse definitions of the form s3://bucket/key
// and https://s3-us-west-2.amazonaws.com/bucket/key
// into bucket and key.
func parseS3ObjectString(obj string) (string, string, error) {
if strings.HasPrefix(obj, prefixS3) {
return parseS3Key(obj)
}
return parseHTTPKey(obj)
}
func parseS3Key(obj string) (string, string, error) {
str := strings.Replace(obj, prefixS3, "", 1)
split := strings.SplitN(str, "/", 2)
if len(split) != 2 {
return "", "", fmt.Errorf("could not parse s3 key: %s", obj)
}
return split[0], split[1], nil
}
func parseHTTPKey(obj string) (string, string, error) {
str := strings.Replace(obj, prefixHTTPS, "", 1)
escaped, err := url.PathUnescape(str)
if err != nil {
return "", "", fmt.Errorf("could not unescape url: %s", obj)
}
split := strings.SplitN(escaped, "/", 3)
if len(split) != 3 {
return "", "", fmt.Errorf("could not parse http s3 key: %s", obj)
}
return split[1], split[2], nil
}
func validateCredentials() error {
sess := session.Must(session.NewSession())
creds := sess.Config.Credentials
_, err := creds.Get()
return err
}