Allow date templates for file names

input_file_test.go

@@ -161,7 +161,7 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
 		requestGenerator.wg.Done()
 	})
 
-	outputFile := NewFileOutput(f.Name())
+	outputFile := NewFileOutput(f.Name(), time.Minute)
 
 	Plugins.Inputs = requestGenerator.inputs
 	Plugins.Outputs = []io.Writer{output, outputFile}

output_file.go

@@ -1,34 +1,70 @@
 package main
 
 import (
-	"io"
+	"bufio"
+	"fmt"
 	"log"
 	"os"
+	"strings"
+	"time"
 )
 
+var dateFileNameFuncs = map[string]func() string{
+	"%Y":  func() string { return time.Now().Format("2006") },
+	"%m":  func() string { return time.Now().Format("01") },
+	"%d":  func() string { return time.Now().Format("02") },
+	"%H":  func() string { return time.Now().Format("15") },
+	"%M":  func() string { return time.Now().Format("04") },
+	"%S":  func() string { return time.Now().Format("05") },
+	"%NS": func() string { return fmt.Sprint(time.Now().Nanosecond()) },
+}
+
 // FileOutput output plugin
 type FileOutput struct {
-	path string
-	file *os.File
+	pathTemplate string
+	currentName  string
+	file         *os.File
+	writer       *bufio.Writer
 }
 
 // NewFileOutput constructor for FileOutput, accepts path
-func NewFileOutput(path string) io.Writer {
+func NewFileOutput(pathTemplate string, flushInterval time.Duration) *FileOutput {
 	o := new(FileOutput)
-	o.path = path
-	o.init(path)
+	o.pathTemplate = pathTemplate
+	o.updateName()
+
+	// Force flushing every minute
+	go func() {
+		for {
+			time.Sleep(flushInterval)
+			if err := o.writer.Flush(); err != nil {
+				break
+			}
+		}
+	}()
+
+	go func() {
+		for {
+			time.Sleep(time.Second)
+			o.updateName()
+		}
+	}()
 
 	return o
 }
 
-func (o *FileOutput) init(path string) {
-	var err error
+func (o *FileOutput) filename() string {
+	path := o.pathTemplate
 
-	o.file, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
-
-	if err != nil {
-		log.Fatal(o, "Cannot open file %q. Error: %s", path, err)
+	for name, fn := range dateFileNameFuncs {
+		path = strings.Replace(path, name, fn(), -1)
 	}
+
+	return path
+}
+
+func (o *FileOutput) updateName() {
+	o.currentName = o.filename()
 }
 
 func (o *FileOutput) Write(data []byte) (n int, err error) {
@@ -36,12 +72,37 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
 		return len(data), nil
 	}
 
-	o.file.Write(data)
-	o.file.Write([]byte(payloadSeparator))
+	if o.file == nil || o.currentName != o.file.Name() {
+		if o.file != nil {
+			o.writer.Flush()
+			o.file.Close()
+		}
+
+		o.file, err = os.OpenFile(o.currentName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
+		o.writer = bufio.NewWriter(o.file)
+
+		if err != nil {
+			log.Fatal(o, "Cannot open file %q. Error: %s", o.currentName, err)
+		}
+	}
+
+	o.writer.Write(data)
+	o.writer.Write([]byte(payloadSeparator))
 
 	return len(data), nil
 }
 
+func (o *FileOutput) Flush() {
+	if o.file != nil {
+		o.writer.Flush()
+	}
+}
+
 func (o *FileOutput) String() string {
-	return "File output: " + o.path
+	return "File output: " + o.file.Name()
+}
+
+func (o *FileOutput) Close() {
+	o.writer.Flush()
+	o.file.Close()
 }

output_file_test.go

@@ -1,9 +1,14 @@
 package main
 
 import (
+	"fmt"
 	"io"
+	"log"
+	"os"
 	"sync"
+	"sync/atomic"
 	"testing"
+	"time"
 )
 
 func TestFileOutput(t *testing.T) {
@@ -11,7 +16,7 @@ func TestFileOutput(t *testing.T) {
 	quit := make(chan int)
 
 	input := NewTestInput()
-	output := NewFileOutput("/tmp/test_requests.gor")
+	output := NewFileOutput("/tmp/test_requests.gor", time.Minute)
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -23,12 +28,18 @@ func TestFileOutput(t *testing.T) {
 		input.EmitGET()
 		input.EmitPOST()
 	}
+	time.Sleep(100 * time.Millisecond)
+	output.Flush()
+
 	close(quit)
 
 	quit = make(chan int)
 
+	var counter int64
 	input2 := NewFileInput("/tmp/test_requests.gor")
 	output2 := NewTestOutput(func(data []byte) {
+		atomic.AddInt64(&counter, 1)
+		log.Println(counter)
 		wg.Done()
 	})
 
@@ -40,3 +51,44 @@ func TestFileOutput(t *testing.T) {
 	wg.Wait()
 	close(quit)
 }
+
+func TestFileOutputPathTemplate(t *testing.T) {
+	output := &FileOutput{pathTemplate: "/tmp/log-%Y-%m-%d-%S"}
+	now := time.Now()
+	expectedPath := fmt.Sprintf("/tmp/log-%s-%s-%s-%s", now.Format("2006"), now.Format("01"), now.Format("02"), now.Format("05"))
+	path := output.filename()
+
+	if expectedPath != path {
+		t.Errorf("Expected path %s but got %s", expectedPath, path)
+	}
+}
+
+func TestFileOutputMultipleFiles(t *testing.T) {
+	output := NewFileOutput("/tmp/log-%Y-%m-%d-%S", time.Minute)
+
+	if output.file != nil {
+		t.Error("Should not initialize file if no writes")
+	}
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name1 := output.file.Name()
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name2 := output.file.Name()
+
+	time.Sleep(time.Second)
+
+	output.Write([]byte("1 1 1\r\ntest"))
+	name3 := output.file.Name()
+
+	if name2 != name1 {
+		t.Errorf("Fast changes should happen in same file:", name1, name2)
+	}
+
+	if name3 == name1 {
+		t.Errorf("File name should change:", name1, name3)
+	}
+
+	os.Remove(name1)
+	os.Remove(name3)
+}

plugins.go

@@ -117,7 +117,7 @@ func InitPlugins() {
 	}
 
 	for _, options := range Settings.outputFile {
-		registerPlugin(NewFileOutput, options)
+		registerPlugin(NewFileOutput, options, Settings.outputFileFlushInterval)
 	}
 
 	for _, options := range Settings.inputHTTP {

settings.go

@@ -39,8 +39,9 @@ type AppSettings struct {
 	outputTCP      MultiOption
 	outputTCPStats bool
 
-	inputFile  MultiOption
-	outputFile MultiOption
+	inputFile               MultiOption
+	outputFile              MultiOption
+	outputFileFlushInterval time.Duration
 
 	inputRAW              MultiOption
 	inputRAWEngine        string
@@ -84,6 +85,7 @@ func init() {
 
 	flag.Var(&Settings.inputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com")
 	flag.Var(&Settings.outputFile, "output-file", "Write incoming requests to file: \n\tgor --input-raw :80 --output-file ./requests.gor")
+	flag.DurationVar(&Settings.outputFileFlushInterval, "output-file-flush-interval", time.Minute, "Interval for forcing buffer flush to the file, default: 60s.")
 
 	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")