mirror of
https://github.com/buger/goreplay.git
synced 2024-04-21 12:32:02 +00:00
s3-output
This commit is contained in:
Leonid Bugaev
@@ -14,6 +14,7 @@ RUN apt-get install flex bison -y
|
||||
RUN wget http://www.tcpdump.org/release/libpcap-1.7.4.tar.gz && tar xzf libpcap-1.7.4.tar.gz && cd libpcap-1.7.4 && ./configure && make install
|
||||
RUN go get github.com/google/gopacket
|
||||
RUN go get -u github.com/golang/lint/golint
|
||||
RUN go get -u github.com/aws/aws-sdk-go
|
||||
|
||||
WORKDIR /go/src/github.com/buger/gor/
|
||||
ADD . /go/src/github.com/buger/gor/
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go middleware.go protocol.go output_file_settings.go
|
||||
SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go middleware.go protocol.go output_file_settings.go output_s3.go
|
||||
SOURCE_PATH = /go/src/github.com/buger/gor/
|
||||
PORT = 8000
|
||||
FADDR = :8000
|
||||
|
||||
+185
@@ -0,0 +1,185 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
_ "bufio"
|
||||
"fmt"
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"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"
|
||||
"io"
|
||||
"log"
|
||||
"math/rand"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sort"
|
||||
"crypto/sha1"
|
||||
"encoding/hex"
|
||||
)
|
||||
|
||||
type S3InputConfig struct {
|
||||
bufferConfig FileInputConfig
|
||||
|
||||
bufferPath string
|
||||
region string
|
||||
endpoint string
|
||||
}
|
||||
|
||||
// FileOutput output plugin
|
||||
type S3Output struct {
|
||||
pathTemplate string
|
||||
|
||||
buffer *FileInput
|
||||
session *session.Session
|
||||
config *S3InputConfig
|
||||
}
|
||||
|
||||
// NewFileOutput constructor for FileOutput, accepts path
|
||||
func NewS3Input(pathTemplate string, config *S3InputConfig) *S3Input {
|
||||
o := new(S3Input)
|
||||
o.pathTemplate = pathTemplate
|
||||
o.config = config
|
||||
|
||||
if config.region == "" {
|
||||
config.region = "us-east-1"
|
||||
}
|
||||
|
||||
if config.bufferPath == "" {
|
||||
config.bufferPath = "/tmp"
|
||||
}
|
||||
|
||||
o.connect()
|
||||
|
||||
if !strings.HasPrefix(pathTemplate, "s3://") {
|
||||
log.Fatal("S3 path format should be: s3://<bucket>/<path_format>")
|
||||
}
|
||||
|
||||
return o
|
||||
}
|
||||
|
||||
func (o *S3Output) connect() {
|
||||
if o.session == nil {
|
||||
o.session = session.New(&aws.Config{Region: aws.String(o.config.region)})
|
||||
}
|
||||
}
|
||||
|
||||
type sortByS3FileIndex []*s3.Object
|
||||
|
||||
func (s sortByS3FileIndex) Len() int {
|
||||
return len(s)
|
||||
}
|
||||
|
||||
func (s sortByS3FileIndex) Swap(i, j int) {
|
||||
s[i], s[j] = s[j], s[i]
|
||||
}
|
||||
|
||||
func (s sortByS3FileIndex) Less(i, j int) bool {
|
||||
if withoutIndex(*s[i].Key) == withoutIndex(*s[j].Key) {
|
||||
return getFileIndex(*s[i].Key) < getFileIndex(*s[j].Key)
|
||||
}
|
||||
|
||||
return s[i] < s[j]
|
||||
}
|
||||
|
||||
func (o *S3Input) updateBuffer() (err error) {
|
||||
path := o.pathTemplate[5:] // stripping `s3://`
|
||||
sep := strings.IndexByte(path, '/')
|
||||
|
||||
bucket = path[:sep]
|
||||
key = path[sep+1:]
|
||||
|
||||
params := &s3.ListObjectsInput{
|
||||
Bucket: bucket,
|
||||
Prefix: key,
|
||||
}
|
||||
resp, err := svc.ListObjects(params)
|
||||
sort.Sort(sortByS3FileIndex(resp.Contents))
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if o.buffer.currentFile == nil {
|
||||
fileToDownload := resp.Contents[0]
|
||||
} else {
|
||||
found := false
|
||||
|
||||
bufName := filepath.Base(i.buffer.currentFile.Name())
|
||||
bufHex := strings.TrimSuffix(bufName, filepath.Ext(bufName))
|
||||
bufSha, _ := hex.DecodeString(bufHex)
|
||||
|
||||
for idx, c := range resp.Contents {
|
||||
sha := sha1.Sum(*c.Key)[0:]
|
||||
|
||||
if bytes.Equal(bufSha, sha) && idx != len(matches)-1 {
|
||||
if i.buffer.currentFile, err = os.Open(matches[idx+1]); err != nil {
|
||||
log.Println("Can't read file ", matches[idx+1], err)
|
||||
return
|
||||
}
|
||||
|
||||
found = true
|
||||
}
|
||||
}
|
||||
|
||||
if !found {
|
||||
return new(NextFileNotFound)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
if strings.HasSuffix(o.pathTemplate, ".gz") {
|
||||
buffer_name += ".gz"
|
||||
}
|
||||
|
||||
buffer_path := filepath.Join(o.config.bufferPath, buffer_name)
|
||||
}
|
||||
|
||||
func (o *S3Output) Read(data []byte) (int, error) {
|
||||
return o.buffer.Read(data)
|
||||
}
|
||||
|
||||
func (o *S3Output) String() string {
|
||||
return "S3 Input: " + o.file.Name()
|
||||
}
|
||||
|
||||
func (o *S3Output) Close() {
|
||||
o.buffer.Close()
|
||||
}
|
||||
|
||||
func (o *S3Output) keyPath(idx int) (bucket, key string) {
|
||||
path := o.pathTemplate[5:] // stripping `s3://`
|
||||
sep := strings.IndexByte(path, '/')
|
||||
|
||||
bucket = path[:sep]
|
||||
key = path[sep+1:]
|
||||
|
||||
for name, fn := range dateFileNameFuncs {
|
||||
key = strings.Replace(key, name, fn(), -1)
|
||||
}
|
||||
|
||||
key = setFileIndex(key, idx)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (o *S3Output) onBufferUpdate(path string) {
|
||||
svc := s3.New(o.session)
|
||||
idx := getFileIndex(path)
|
||||
bucket, key := o.keyPath(idx)
|
||||
|
||||
file, _ := os.Open(path)
|
||||
// reader := bufio.NewReader(file)
|
||||
|
||||
_, err := svc.PutObject(&s3.PutObjectInput{
|
||||
Body: file,
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
log.Printf("Failed to upload data to %s/%s, %s\n", bucket, key, err)
|
||||
return
|
||||
}
|
||||
|
||||
os.Remove(path)
|
||||
}
|
||||
+13
-1
@@ -29,6 +29,7 @@ type FileOutputConfig struct {
|
||||
sizeLimit unitSizeVar
|
||||
queueLimit int
|
||||
append bool
|
||||
onClose func(string)
|
||||
}
|
||||
|
||||
// FileOutput output plugin
|
||||
@@ -50,6 +51,10 @@ func NewFileOutput(pathTemplate string, config *FileOutputConfig) *FileOutput {
|
||||
o.config = config
|
||||
o.updateName()
|
||||
|
||||
if o.config.flushInterval == 0 {
|
||||
o.config.flushInterval = time.Minute
|
||||
}
|
||||
|
||||
// Force flushing every minute
|
||||
go func() {
|
||||
for {
|
||||
@@ -87,7 +92,10 @@ func setFileIndex(name string, idx int) string {
|
||||
withoutExt := strings.TrimSuffix(name, ext)
|
||||
|
||||
if i := strings.LastIndex(withoutExt, "_"); i != -1 {
|
||||
withoutExt = withoutExt[:i]
|
||||
// Only prefixes with numbers counts
|
||||
if _, err := strconv.Atoi(withoutExt[i+1:]); err == nil {
|
||||
withoutExt = withoutExt[:i]
|
||||
}
|
||||
}
|
||||
|
||||
return withoutExt + "_" + idxS + ext
|
||||
@@ -224,5 +232,9 @@ func (o *FileOutput) Close() {
|
||||
o.writer.(*bufio.Writer).Flush()
|
||||
}
|
||||
o.file.Close()
|
||||
|
||||
if o.config.onClose != nil {
|
||||
go o.config.onClose(o.file.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+122
@@ -0,0 +1,122 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
_ "bufio"
|
||||
"fmt"
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"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"
|
||||
"io"
|
||||
"log"
|
||||
"math/rand"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type S3OutputConfig struct {
|
||||
bufferConfig FileOutputConfig
|
||||
|
||||
bufferPath string
|
||||
region string
|
||||
endpoint string
|
||||
}
|
||||
|
||||
// FileOutput output plugin
|
||||
type S3Output struct {
|
||||
pathTemplate string
|
||||
|
||||
buffer *FileOutput
|
||||
session *session.Session
|
||||
config *S3OutputConfig
|
||||
}
|
||||
|
||||
// NewFileOutput constructor for FileOutput, accepts path
|
||||
func NewS3Output(pathTemplate string, config *S3OutputConfig) *S3Output {
|
||||
o := new(S3Output)
|
||||
o.pathTemplate = pathTemplate
|
||||
o.config = config
|
||||
o.config.bufferConfig.onClose = o.onBufferUpdate
|
||||
|
||||
if config.region == "" {
|
||||
config.region = "us-east-1"
|
||||
}
|
||||
|
||||
if config.bufferPath == "" {
|
||||
config.bufferPath = "/tmp"
|
||||
}
|
||||
|
||||
rnd := rand.Int63()
|
||||
buffer_name := fmt.Sprintf("gor_output_s3_%d_buf", rnd)
|
||||
|
||||
if strings.HasSuffix(o.pathTemplate, ".gz") {
|
||||
buffer_name += ".gz"
|
||||
}
|
||||
|
||||
buffer_path := filepath.Join(config.bufferPath, buffer_name)
|
||||
|
||||
o.buffer = NewFileOutput(buffer_path, &config.bufferConfig)
|
||||
o.connect()
|
||||
|
||||
if !strings.HasPrefix(pathTemplate, "s3://") {
|
||||
log.Fatal("S3 path format should be: s3://<bucket>/<path_format>")
|
||||
}
|
||||
|
||||
return o
|
||||
}
|
||||
|
||||
func (o *S3Output) connect() {
|
||||
if o.session == nil {
|
||||
o.session = session.New(&aws.Config{Region: aws.String(o.config.region)})
|
||||
}
|
||||
}
|
||||
|
||||
func (o *S3Output) Write(data []byte) (n int, err error) {
|
||||
return o.buffer.Write(data)
|
||||
}
|
||||
|
||||
func (o *S3Output) String() string {
|
||||
return "File output: " + o.file.Name()
|
||||
}
|
||||
|
||||
func (o *S3Output) Close() {
|
||||
o.buffer.Close()
|
||||
}
|
||||
|
||||
func (o *S3Output) keyPath(idx int) (bucket, key string) {
|
||||
path := o.pathTemplate[5:] // stripping `s3://`
|
||||
sep := strings.IndexByte(path, '/')
|
||||
|
||||
bucket = path[:sep]
|
||||
key = path[sep+1:]
|
||||
|
||||
for name, fn := range dateFileNameFuncs {
|
||||
key = strings.Replace(key, name, fn(), -1)
|
||||
}
|
||||
|
||||
key = setFileIndex(key, idx)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (o *S3Output) onBufferUpdate(path string) {
|
||||
svc := s3.New(o.session)
|
||||
idx := getFileIndex(path)
|
||||
bucket, key := o.keyPath(idx)
|
||||
|
||||
file, _ := os.Open(path)
|
||||
// reader := bufio.NewReader(file)
|
||||
|
||||
_, err := svc.PutObject(&s3.PutObjectInput{
|
||||
Body: file,
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
log.Printf("Failed to upload data to %s/%s, %s\n", bucket, key, err)
|
||||
return
|
||||
}
|
||||
|
||||
os.Remove(path)
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/s3"
|
||||
"math/rand"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestS3Output(t *testing.T) {
|
||||
bucket := aws.String("test-gor")
|
||||
rnd := rand.Int63()
|
||||
path := fmt.Sprintf("s3://test-gor/%d/requests.gz", rnd)
|
||||
|
||||
output := NewS3Output(path,
|
||||
&S3OutputConfig{
|
||||
bufferConfig: FileOutputConfig{queueLimit: 2},
|
||||
},
|
||||
)
|
||||
|
||||
svc := s3.New(output.session)
|
||||
|
||||
output.Write([]byte("1 1 1\ntest"))
|
||||
output.Write([]byte("1 1 1\ntest"))
|
||||
output.buffer.updateName()
|
||||
output.Write([]byte("1 1 1\ntest"))
|
||||
output.Write([]byte("1 1 1\ntest"))
|
||||
output.buffer.updateName()
|
||||
output.Write([]byte("1 1 1\ntest"))
|
||||
|
||||
params := &s3.ListObjectsInput{
|
||||
Bucket: bucket,
|
||||
Prefix: aws.String(fmt.Sprintf("%d", rnd)),
|
||||
}
|
||||
|
||||
resp, _ := svc.ListObjects(params)
|
||||
if len(resp.Contents) != 2 {
|
||||
t.Error("Should create 2 objects", len(resp.Contents))
|
||||
} else {
|
||||
if *resp.Contents[0].Key != fmt.Sprintf("%d/requests_0.gz", rnd) ||
|
||||
*resp.Contents[1].Key != fmt.Sprintf("%d/requests_1.gz", rnd) {
|
||||
t.Error("Should assign proper names", resp.Contents)
|
||||
}
|
||||
}
|
||||
|
||||
for _, c := range resp.Contents {
|
||||
svc.DeleteObject(&s3.DeleteObjectInput{Bucket: bucket, Key: c.Key})
|
||||
}
|
||||
|
||||
matches, _ := filepath.Glob(fmt.Sprintf("/tmp/gor_output_s3_*"))
|
||||
for _, m := range matches {
|
||||
os.Remove(m)
|
||||
}
|
||||
}
|
||||
@@ -120,6 +120,11 @@ func InitPlugins() {
|
||||
registerPlugin(NewFileOutput, options, Settings.outputFileConfig)
|
||||
}
|
||||
|
||||
for _, options := range Settings.outputS3 {
|
||||
Settings.outputS3Config.bufferConfig = Settings.outputFileConfig
|
||||
registerPlugin(NewS3Output, options, Settings.outputS3Config)
|
||||
}
|
||||
|
||||
for _, options := range Settings.inputHTTP {
|
||||
registerPlugin(NewHTTPInput, options)
|
||||
}
|
||||
|
||||
@@ -44,6 +44,9 @@ type AppSettings struct {
|
||||
outputFile MultiOption
|
||||
outputFileConfig FileOutputConfig
|
||||
|
||||
outputS3 MultiOption
|
||||
outputS3Config S3OutputConfig
|
||||
|
||||
inputRAW MultiOption
|
||||
inputRAWEngine string
|
||||
inputRAWTrackResponse bool
|
||||
@@ -97,6 +100,11 @@ func init() {
|
||||
flag.Var(&Settings.outputFileConfig.sizeLimit, "output-file-size-limit", "Size of each chunk. Default: 32mb")
|
||||
flag.IntVar(&Settings.outputFileConfig.queueLimit, "output-file-queue-limit", 256, "The length of the chunk queue. Default: 256")
|
||||
|
||||
flag.Var(&Settings.outputS3, "output-s3", "Write incoming requests to S3 path: \n\tgor --input-raw :80 --output-s3 s3://mybucket/logs/%Y-%m-%d.gz")
|
||||
flag.StringVar(&Settings.outputS3Config.bufferPath, "output-s3-buffer-path", "/tmp", "The path prefix of the S3 log buffer files.: \n\tgor --input-raw :80 --output-s3 s3://mybucket/logs/%Y-%m-%d.gz --output-s3-buffer-path /mnt/logs")
|
||||
flag.StringVar(&Settings.outputS3Config.region, "output-s3-region", "us-east-1", "Specify S3 region, default is 'us-east-1': \n\tgor --input-raw :80 --output-s3 s3://mybucket/logs/%Y-%m-%d.gz --output-s3-region us-west-2")
|
||||
flag.StringVar(&Settings.outputS3Config.endpoint, "output-s3-endpoint", "", "You can specify custom endpoint if using alternative S3 compatitable storage.")
|
||||
|
||||
flag.Var(&Settings.inputRAW, "input-raw", "Capture traffic from given port (use RAW sockets and require *sudo* access):\n\t# Capture traffic from 8080 port\n\tgor --input-raw :8080 --output-http staging.com")
|
||||
|
||||
flag.BoolVar(&Settings.inputRAWTrackResponse, "input-raw-track-response", false, "If turned on Gor will track responses in addition to requests, and they will be available to middleware and file output.")
|
||||
|
||||
Reference in New Issue
Block a user