<|start_filename|>pkg/backend/stream/client_test.go<|end_filename|>
package stream

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/egnyte/ax/pkg/backend/common"
)

func TestFindTimestamp(t *testing.T) {
	sampleData := `{"jstimestamp":1504516581620, "message": "Sup yo"}
{"ts": "2017-08-04T11:16:52.088Z", "message": "Sup yo 2"}
`
	sc := New(strings.NewReader(sampleData))
	for msg := range sc.Query(context.Background(), common.Query{}) {
		//fmt.Printf("%+v\n", msg)
		if msg.Timestamp.Day() != 4 {
			t.Error("Wrong day", msg.Timestamp.Day())
		}
		if msg.Timestamp.Minute() != 16 {
			t.Error("Wrong minute", msg.Timestamp.Minute())
		}
	}
}

func TestFindTimestampInMessage(t *testing.T) {
	sampleData := `{"message": "INFO: This happened at 2017-07-04T14:09:28.184Z so"}
{"message": "2017-06-04 06:52:14,689  INFO    All ELCC processes are running"}
{"message": "[2017-05-04 08:51:14,689]  INFO    All ELCC processes are running"}
{"message": "(2017-06-04 09:25:39,261) INFO    (Processor) End of sync notification sent to server"}
`
	months := []time.Month{7, 6, 5, 6}
	sc := New(strings.NewReader(sampleData))
	counter := 0
	for msg := range sc.Query(context.Background(), common.Query{}) {
		if msg.Timestamp.Month() != months[counter] {
			t.Error("Wrong month", msg.Timestamp.Month(), "expected", months[counter], "in", msg.Attributes["message"])
		}
		message, _ := msg.Attributes["message"].(string)
		if !strings.HasPrefix(message, "INFO") {
			t.Error("Stripping of timestamp failed")
		}

		counter++
	}

}


<|start_filename|>pkg/alert/alert.go<|end_filename|>
package alert

import "github.com/egnyte/ax/pkg/backend/common"

type Alerter interface {
	SendAlert(lm common.LogMessage) error
}


<|start_filename|>pkg/backend/kibana/util.go<|end_filename|>
package kibana

import (
	"bytes"
	"encoding/json"
	"io"
	"regexp"
	"time"
)

func createMultiSearch(objs ...interface{}) (io.Reader, error) {
	var buf bytes.Buffer
	encoder := json.NewEncoder(&buf)
	for _, obj := range objs {
		err := encoder.Encode(obj)
		if err != nil {
			return nil, err
		}
	}
	//fmt.Println(buf.String())
	return &buf, nil
}

func safeFilename(name string) string {
	re := regexp.MustCompile(`[^\w\-]`)
	return re.ReplaceAllString(name, "_")
}

func unixMillis(t time.Time) int64 {
	return t.Unix() * 1000
}


<|start_filename|>pkg/cache/cache.go<|end_filename|>
package cache

import (
	"encoding/json"
	"log"
	"os"
	"time"
)

type cacheItem struct {
	Value      interface{} `json:"value"`
	ExpireDate *time.Time  `json:"expire"`
}

type Cache struct {
	path string
	data map[string]cacheItem
}

func New(path string) *Cache {
	cache := &Cache{
		path: path,
		data: make(map[string]cacheItem),
	}
	file, err := os.Open(cache.path)
	if err == nil {
		defer file.Close()
		decoder := json.NewDecoder(file)
		err = decoder.Decode(&cache.data)
		if err != nil {
			log.Printf("Error decoding cache: %v\n", err)
			// Let's recover gracefully
			cache.data = make(map[string]cacheItem)
			cache.Flush()
		}
	}
	return cache
}

func (cache *Cache) Contains(key string) bool {
	if value, ok := cache.data[key]; ok {
		if value.ExpireDate != nil && value.ExpireDate.Before(time.Now()) {
			delete(cache.data, key)
			return false
		}
		return true
	}
	return false
}

func (cache *Cache) Set(key string, value interface{}, expire *time.Time) {
	cache.data[key] = cacheItem{value, expire}
}

func (cache *Cache) Unset(key string) {
	delete(cache.data, key)
}

func (cache *Cache) Get(key string) interface{} {
	// Could be optimized slightly, but oh well
	if cache.Contains(key) {
		return cache.data[key].Value
	} else {
		return nil
	}
}

func (cache *Cache) GetString(key string) string {
	return cache.Get(key).(string)
}

func (cache *Cache) GetMap(key string) map[string]interface{} {
	return cache.Get(key).(map[string]interface{})
}

func (cache *Cache) Flush() error {
	file, err := os.OpenFile(cache.path, os.O_CREATE|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	file.Truncate(0)
	defer file.Close()
	encoder := json.NewEncoder(file)
	return encoder.Encode(cache.data)
}

func (cache *Cache) Remove() error {
	return os.Remove(cache.path)
}


<|start_filename|>cmd/ax/alert.go<|end_filename|>
package main

import (
	"context"
	"fmt"
	"time"

	"github.com/egnyte/ax/pkg/alert"
	"github.com/egnyte/ax/pkg/alert/slack"
	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/config"
)

var (
	alertFlags    = addQueryFlags(addAlertCommand)
	alertFlagName string
)

func init() {
	addAlertCommand.Flag("name", "Name for alert").Required().StringVar(&alertFlagName)
}

func addAlertMain(rc config.RuntimeConfig, client common.Client) {
	alertConfig := config.AlertConfig{
		Env:      rc.ActiveEnv,
		Name:     alertFlagName,
		Selector: *alertFlags,
	}

	fmt.Printf("Config: %+v\n", alertConfig)
	conf := config.LoadConfig()
	conf.Alerts = append(conf.Alerts, alertConfig)
	config.SaveConfig(conf)
}

func watchAlerts(ctx context.Context, rc config.RuntimeConfig, alertConfig config.AlertConfig) {
	var alerter alert.Alerter
	switch alertConfig.Service["backend"] {
	case "slack":
		alerter = slack.New(alertConfig.Name, rc.DataDir, alertConfig.Service)
	default:
		fmt.Printf("Back-end type not supported: %s\n", alertConfig.Service["backend"])
		return
	}
	query := querySelectorsToQuery(&alertConfig.Selector)
	query.Follow = true
	query.MaxResults = 100
	client := determineClient(rc.Config.Environments[alertConfig.Env])
	if client == nil {
		fmt.Println("Cannot obtain a client for", alertConfig)
		return
	}
	fmt.Println("Now waiting for alerts for", alertConfig.Name)
	for message := range client.Query(ctx, query) {
		fmt.Printf("[%s] Sending alert to %s: %+v\n", alertConfig.Name, alertConfig.Service["backend"], message.Map())
		err := alerter.SendAlert(message)
		if err != nil {
			fmt.Println("Couldn't send alert", err)
			continue
		}
	}
}

func alertMain(ctx context.Context, rc config.RuntimeConfig) {
	for _, alert := range rc.Config.Alerts {
		go watchAlerts(ctx, rc, alert)
	}
	for {
		time.Sleep(time.Minute)
	}
}


<|start_filename|>pkg/backend/common/client_test.go<|end_filename|>
package common

import (
	"testing"
	"time"
)

func TestFilter(t *testing.T) {
	lm := LogMessage{
		Timestamp: time.Now(),
		Attributes: map[string]interface{}{
			"message": "Sup",
			"someStr": "Zef",
			"someN":   34,
		},
	}
	lastHour := time.Now().Add(-time.Hour)
	nextHour := time.Now().Add(time.Hour)
	shouldMatchQueries := []Query{
		{
			EqualityFilters: []EqualityFilter{
				{FieldName: "someStr", Value: "Zef", Operator: "="},
			},
		},
		{
			EqualityFilters: []EqualityFilter{
				{FieldName: "someN", Value: "34", Operator: "="},
			},
		},
		{
			QueryString: "zef",
			EqualityFilters: []EqualityFilter{
				{FieldName: "someN", Value: "34", Operator: "="},
			},
		},
		{
			QueryString: "zef",
			EqualityFilters: []EqualityFilter{
				{FieldName: "someN", Value: "34", Operator: "="},
			},
			Before: &nextHour,
			After:  &lastHour,
		},
		{
			EqualityFilters: []EqualityFilter{
				{FieldName: "someN", Value: "32", Operator: "!="},
			},
		},
		{
			EqualityFilters: []EqualityFilter{
				{FieldName: "someNonexistingField", Value: "Pete", Operator: "!="},
			},
		},
		{
			ExistenceFilters: []ExistenceFilter{
				{
					FieldName: "message",
					Exists:    true,
				},
			},
		},
		{
			ExistenceFilters: []ExistenceFilter{
				{
					FieldName: "message",
					Exists:    true,
				},
				{
					FieldName: "someStr",
					Exists:    true,
				},
			},
		},
		{
			ExistenceFilters: []ExistenceFilter{
				{
					FieldName: "message",
					Exists:    true,
				},
				{
					FieldName: "someStr",
					Exists:    true,
				},
				{
					FieldName: "bar",
					Exists:    false,
				},
			},
		},
		{
			MembershipFilters: []MembershipFilter{
				{
					FieldName:   "message",
					ValidValues: []string{"Sup", "bar"},
				},
			},
		},
		{
			MembershipFilters: []MembershipFilter{
				{
					FieldName:     "message",
					InvalidValues: []string{"foo", "bar"},
				},
			},
		},
		{
			MembershipFilters: []MembershipFilter{
				{
					FieldName:   "someN",
					ValidValues: []string{"Sup", "34"},
				},
			},
		},
	}
	shouldNotMatchQueries := []Query{
		{
			EqualityFilters: []EqualityFilter{
				{FieldName: "someStr", Value: "Pete", Operator: "="},
			},
		},
		{
			QueryString: "bla",
			EqualityFilters: []EqualityFilter{
				{FieldName: "someStr", Value: "Pete", Operator: "="},
			},
		},
		{
			After: &nextHour,
		},
		{
			ExistenceFilters: []ExistenceFilter{
				{
					FieldName: "message",
					Exists:    true,
				},
				{
					FieldName: "someStr",
					Exists:    false,
				},
				{
					FieldName: "bar",
					Exists:    false,
				},
			},
		},
		{
			ExistenceFilters: []ExistenceFilter{
				{
					FieldName: "bar",
					Exists:    true,
				},
			},
		},
		{
			MembershipFilters: []MembershipFilter{
				{
					FieldName:   "someN",
					ValidValues: []string{"45", "bar"},
				},
			},
		},
		{
			MembershipFilters: []MembershipFilter{
				{
					FieldName:     "someN",
					InvalidValues: []string{"34"},
				},
			},
		},
	}
	for i, shouldMatch := range shouldMatchQueries {
		if !MatchesQuery(lm, shouldMatch) {
			t.Errorf("Did not match: %d: %+v", i, shouldMatch)
		}
	}
	for i, shouldNotMatch := range shouldNotMatchQueries {
		if MatchesQuery(lm, shouldNotMatch) {
			t.Errorf("Did match: %d: %+v", i, shouldNotMatch)
		}
	}
}

func TestFlatten(t *testing.T) {
	into := make(map[string]interface{})
	inputJsonString := `{
		"name": "<NAME>",
		"age": 34,
		"docker": {
			"service": "myservice",
			"container": "something else",
			"deeper": {
				"b": 10
			}
		}
	}`
	expectedJsonString := `{
		"name": "<NAME>",
		"age": 34,
		"docker.service": "myservice",
		"docker.container": "something else",
		"docker.deeper.b": 10
	}`
	var inputObj map[string]interface{}
	var expectedObj map[string]interface{}
	MustJsonDecode(inputJsonString, &inputObj)
	MustJsonDecode(expectedJsonString, &expectedObj)
	FlattenAttributes(inputObj, into, "")

	if MustJsonEncode(into) != MustJsonEncode(expectedObj) {
		t.Error("Didn't match")
	}
}

func TestMembershipFilter_Matches(t *testing.T) {
	type fields struct {
		FieldName     string
		ValidValues   []string
		InvalidValues []string
	}
	tests := []struct {
		name   string
		fields fields
		m      LogMessage
		want   bool
	}{
		{
			name: "Enforce field existence for inclusive member filters",
			fields: fields{
				FieldName: "domain",
				ValidValues: []string{
					"ax",
				},
				InvalidValues: []string{},
			},
			m: LogMessage{
				Attributes: map[string]interface{}{
					"not-domain": "foo",
				},
			},
			want: false,
		},
		{
			name: "Don't enforce field existence for exclusive-only filter",
			fields: fields{
				FieldName:   "domain",
				ValidValues: []string{},
				InvalidValues: []string{
					"ax",
				},
			},
			m: LogMessage{
				Attributes: map[string]interface{}{
					"not-domain": "foo",
				},
			},
			want: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			f := MembershipFilter{
				FieldName:     tt.fields.FieldName,
				ValidValues:   tt.fields.ValidValues,
				InvalidValues: tt.fields.InvalidValues,
			}
			if got := f.Matches(tt.m); got != tt.want {
				t.Errorf("MembershipFilter.Matches() = %v, want %v", got, tt.want)
			}
		})
	}
}


<|start_filename|>pkg/config/cloudwatch.go<|end_filename|>
package config

import (
	"bufio"
	"fmt"

	"github.com/egnyte/ax/pkg/backend/cloudwatch"
)

func cloudwatchConfig(reader *bufio.Reader, existingConfig Config) (EnvMap, error) {
	em := EnvMap{
		"backend": "cloudwatch",
	}
	existingCwEnv := findFirstEnvWhere(existingConfig.Environments, func(em EnvMap) bool {
		return em["backend"] == "cloudwatch"
	})
	if existingCwEnv != nil {
		accessKey := (*existingCwEnv)["accesskey"]
		fmt.Printf("Access Key ID [%s]: ", accessKey)
		em["accesskey"] = readLine(reader)
		if em["accesskey"] == "" {
			em["accesskey"] = (*existingCwEnv)["accesskey"]
			em["accesssecretkey"] = (*existingCwEnv)["accesssecretkey"]
		}
	} else {
		fmt.Print("Access Secret Key: ")
		em["accesssecretkey"] = readLine(reader)
	}
	fmt.Print("Region (us-east-1, us-west-1, us-west-2, eu-west-1, eu-central-1, ap-southeast-1, ap-southeast-2, ap-northeast-1, sa-east-1): ")
	em["region"] = readLine(reader)
	var cwClient *cloudwatch.CloudwatchClient
	var groups []string
	var err error
	fmt.Println("Attempting to connect to Cloudwatch")
	cwClient = cloudwatch.New(em["accesskey"], em["accesssecretkey"], em["region"], "")
	groups, err = cwClient.ListGroups()
	if err != nil {
		fmt.Printf("Got error connecting to Cloudwatch: %s\n", err)
		return em, err
	}
	fmt.Println("List of groups:")
	for _, group := range groups {
		fmt.Println("  ", group)
	}
	fmt.Print("Group: ")
	em["groupname"] = readLine(reader)
	return em, nil
}


<|start_filename|>pkg/backend/cloudwatch/client.go<|end_filename|>
package cloudwatch

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/egnyte/ax/pkg/backend/common"
)

type CloudwatchClient struct {
	logs      *cloudwatchlogs.CloudWatchLogs
	groupName string
}

func attemptParseJSON(str string) map[string]interface{} {
	m := make(map[string]interface{})
	// Find start of JSON blob
	startIdx := strings.Index(str, "{")
	if startIdx == -1 { // If not found, fall back to dumping the whole thing into the "message" field
		m["message"] = str
		return m
	}
	err := json.Unmarshal([]byte(str[startIdx:]), &m)
	if err != nil {
		m["message"] = str
	}
	return m
}

func logEventToMessage(query common.Query, logEvent *cloudwatchlogs.FilteredLogEvent) common.LogMessage {
	message := common.NewLogMessage()
	message.ID = *logEvent.EventId
	message.Timestamp = time.Unix((*logEvent.Timestamp)/1000, (*logEvent.Timestamp)%1000)
	message.Attributes = common.Project(attemptParseJSON(*logEvent.Message), query.SelectFields)
	return message
}

// https://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/FilterAndPatternSyntax.html
func queryToFilterPattern(query common.Query) string {
	filterParts := make([]string, 0)
	for _, filter := range query.EqualityFilters {
		filterParts = append(filterParts, fmt.Sprintf("($.%s %s \"%s\")", filter.FieldName, filter.Operator, filter.Value))
	}
	var filterPattern string
	if len(query.EqualityFilters) == 0 {
		filterPattern = query.QueryString
	} else {
		filterPattern = fmt.Sprintf("%s { %s }", query.QueryString, strings.Join(filterParts, " && "))
	}

	return strings.TrimSpace(filterPattern)
}

func (client *CloudwatchClient) ImplementsAdvancedFilters() bool {
	return false
}

func (client *CloudwatchClient) readLogBatch(ctx context.Context, query common.Query) ([]common.LogMessage, error) {
	var startTime, endTime *int64 = nil, nil
	if query.After != nil {
		startTimeVal := (*query.After).UnixNano() / int64(time.Millisecond)
		startTime = &startTimeVal
	}
	if query.Before != nil {
		endTimeVal := (*query.Before).UnixNano() / int64(time.Millisecond)
		endTime = &endTimeVal
	}
	resp, err := client.logs.FilterLogEventsWithContext(ctx, &cloudwatchlogs.FilterLogEventsInput{
		LogGroupName:  aws.String(client.groupName),
		FilterPattern: aws.String(queryToFilterPattern(query)),
		Limit:         aws.Int64(int64(query.MaxResults)),
		StartTime:     startTime,
		EndTime:       endTime,
	})
	if err != nil {
		return nil, err
	}
	messages := make([]common.LogMessage, 0, 20)
	for _, message := range resp.Events {
		messages = append(messages, logEventToMessage(query, message))
	}
	return messages, nil
}

func (client *CloudwatchClient) Query(ctx context.Context, query common.Query) <-chan common.LogMessage {
	if query.Follow {
		return common.ReQueryFollow(ctx, func() ([]common.LogMessage, error) {
			return client.readLogBatch(ctx, query)
		})
	}
	resultChan := make(chan common.LogMessage)

	go func() {
		messages, err := client.readLogBatch(ctx, query)
		if err != nil {
			fmt.Printf("Error while fetching logs: %s\n", err)
			close(resultChan)
			return
		}
		for _, message := range messages {
			resultChan <- message
		}
		close(resultChan)
	}()

	return resultChan
}

func (client *CloudwatchClient) ListGroups() ([]string, error) {
	resp, err := client.logs.DescribeLogGroups(&cloudwatchlogs.DescribeLogGroupsInput{})
	if err != nil {
		return nil, err
	}

	groupNames := make([]string, 0)
	for _, stream := range resp.LogGroups {
		groupNames = append(groupNames, *stream.LogGroupName)
	}

	return groupNames, err
}

func New(accessKey, accessSecretKey, region, groupName string) *CloudwatchClient {
	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String(region),
		Credentials: credentials.NewStaticCredentials(accessKey, accessSecretKey, ""),
	})

	if err != nil {
		fmt.Printf("Could not create AWS Session: %s\n", err)
		return nil
	}
	logs := cloudwatchlogs.New(sess)

	return &CloudwatchClient{
		logs:      logs,
		groupName: groupName,
	}

}

var _ common.Client = &CloudwatchClient{}


<|start_filename|>pkg/backend/stream/client.go<|end_filename|>
package stream

import (
	"bufio"
	"context"

	"io"

	"encoding/json"
	"strings"

	"time"

	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/heuristic"
)

type Client struct {
	reader io.Reader
}

func New(file io.Reader) *Client {
	return &Client{file}
}

func parseLine(line string) common.LogMessage {
	decoder := json.NewDecoder(strings.NewReader(line))
	obj := make(map[string]interface{})
	err := decoder.Decode(&obj)
	if err != nil {
		obj["message"] = strings.TrimSpace(line)
		return common.LogMessage{
			Timestamp:  time.Now(),
			Attributes: obj,
		}
	}
	return common.LogMessage{
		Timestamp:  time.Now(), // TODO: Fix this
		Attributes: obj,
	}
}

func (client *Client) ImplementsAdvancedFilters() bool {
	return true
}

func (client *Client) Query(ctx context.Context, q common.Query) <-chan common.LogMessage {
	resultChan := make(chan common.LogMessage)
	reader := bufio.NewReader(client.reader)
	go func() {
		var ltFunc heuristic.LogTimestampParser
	LFor:
		for {
			select {
			case <-ctx.Done():
				// Context canceled, let's get outta here
				break LFor
			default:
			}
			line, err := reader.ReadString('\n')
			if err != nil {
				// TODO: erro != EOF break?
				if err == io.EOF {
					break
				}
				//fmt.Println("Error: ", err)
				break
			}
			message := parseLine(line)
			if ltFunc == nil {
				ltFunc = heuristic.FindTimestampFunc(message)
			}
			if ltFunc != nil {
				ts := ltFunc(message)
				if ts != nil {
					message.Timestamp = *ts
				} else {
					ltFunc = heuristic.FindTimestampFunc(message)
					if ltFunc != nil {
						ts := ltFunc(message)
						message.Timestamp = *ts
					}
				}
			}
			if common.MatchesQuery(message, q) {
				message.Attributes = common.Project(message.Attributes, q.SelectFields)
				resultChan <- message
			}
		}
		close(resultChan)
	}()

	return resultChan
}

var _ common.Client = &Client{}


<|start_filename|>pkg/backend/kibana/client.go<|end_filename|>
package kibana

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"

	"github.com/egnyte/ax/pkg/backend/common"
)

type Client struct {
	URL        string
	AuthHeader string
	Index      string
}

func New(url, authHeader, index string) *Client {
	return &Client{
		URL:        url,
		AuthHeader: authHeader,
		Index:      index,
	}
}

func (client *Client) ImplementsAdvancedFilters() bool {
	return false
}

func (client *Client) addHeaders(req *http.Request) {
	req.Header.Set("Authorization", client.AuthHeader)
	// TODO: This may seem crazy but this header needs to be set, even if empty
	req.Header.Set("Kbn-Version", "")
	req.Header.Set("Content-Type", "application/x-ldjson")
}

type indexList struct {
	SavedObjects []struct {
		Type       string `json:"type"`
		Attributes struct {
			Title string `json:"title"`
		}
	} `json:"saved_objects"`
}

func (client *Client) ListIndices() ([]string, error) {
	req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/saved_objects/?type=index-pattern&per_page=10000", client.URL), nil)
	if err != nil {
		return nil, err
	}
	client.addHeaders(req)
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusUnauthorized {
		return nil, errors.New("Authentication failed")
	} else if resp.StatusCode != http.StatusOK {
		return nil, errors.New(resp.Status)
	}
	decoder := json.NewDecoder(resp.Body)
	var data indexList
	err = decoder.Decode(&data)
	if err != nil {
		return nil, err
	}
	// Build list
	indexNames := make([]string, 0, len(data.SavedObjects))
	for _, indexInfo := range data.SavedObjects {
		if indexInfo.Type == "index-pattern" {
			indexNames = append(indexNames, indexInfo.Attributes.Title)
		}
	}
	return indexNames, nil
}

var _ common.Client = &Client{}


<|start_filename|>cmd/ax/main.go<|end_filename|>
package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/zefhemel/kingpin"

	"github.com/egnyte/ax/pkg/backend/cloudwatch"
	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/backend/docker"
	"github.com/egnyte/ax/pkg/backend/kibana"
	"github.com/egnyte/ax/pkg/backend/stackdriver"
	"github.com/egnyte/ax/pkg/backend/stream"
	"github.com/egnyte/ax/pkg/backend/subprocess"
	"github.com/egnyte/ax/pkg/config"
)

var (
	queryCommand    = kingpin.Command("query", "Query logs").Default()
	alertCommand    = kingpin.Command("alert", "Be alerted when logs match a query")
	alertDCommand   = kingpin.Command("alertd", "Be alerted when logs match a query")
	versionCommand  = kingpin.Command("version", "Show the ax version")
	upgrade         = kingpin.Command("upgrade", "Upgrade Ax if a new version is available")
	addAlertCommand = alertCommand.Command("add", "Add new alert")
	version         = "dev"
	versionFlag     = kingpin.Version(version)
)

func determineClient(em config.EnvMap) common.Client {
	stat, _ := os.Stdin.Stat()
	var client common.Client
	if (stat.Mode() & os.ModeCharDevice) == 0 {
		client = stream.New(os.Stdin)
	} else {
		switch em["backend"] {
		case "docker":
			client = docker.New(em["pattern"])
		case "kibana":
			client = kibana.New(em["url"], em["auth"], em["index"])
		case "cloudwatch":
			client = cloudwatch.New(em["accesskey"], em["accesssecretkey"], em["region"], em["groupname"])
		case "stackdriver":
			client = stackdriver.New(em["credentials"], em["project"], em["log"])
		case "subprocess":
			client = subprocess.New(strings.Split(em["command"], " "))
		}
	}
	return client
}

func sigtermContextHandler(ctx context.Context) context.Context {
	ctx, cancel := context.WithCancel(ctx)

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		switch <-c {
		case os.Interrupt:
			fmt.Println("Canceled through SIGINT (Ctrl-C)")
		case syscall.SIGTERM:
			fmt.Println("Canceled through SIGTERM")
		}
		cancel()
	}()

	return ctx
}

func main() {
	cmd := kingpin.Parse()

	rc := config.BuildConfig()
	client := determineClient(rc.Env)

	switch cmd {
	case "query":
		ctx := sigtermContextHandler(context.Background())
		if client == nil {
			if len(rc.Config.Environments) == 0 {
				// Assuming first time use
				fmt.Println("Welcome to ax! It looks like this is the first time running, so let's start with creating a new environment.")
				config.AddEnv()
				return
			}
			fmt.Println("No default environment set, please use the --env flag to set one. Exiting.")
			return
		}
		queryMain(ctx, rc, client)
	case "env add":
		config.AddEnv()
	case "env list":
		config.ListEnvs()
	case "env edit":
		config.EditConfig()
	case "alert add":
		addAlertMain(rc, client)
	case "alertd":
		alertMain(context.Background(), rc)
	case "version":
		println(version)
	case "upgrade":
		if err := upgradeVersion(); err != nil {
			fmt.Println("Upgrade failed.")
		} else {
			fmt.Println("Upgrade has been completed successfully.")
		}
	}
}


<|start_filename|>pkg/config/kibana.go<|end_filename|>
package config

import (
	"bufio"
	"fmt"

	"github.com/egnyte/ax/pkg/backend/kibana"
)

func kibanaConfig(reader *bufio.Reader, existingConfig Config) (EnvMap, error) {
	em := EnvMap{
		"backend": "kibana",
	}
	existingKibanaEnv := findFirstEnvWhere(existingConfig.Environments, func(em EnvMap) bool {
		return em["backend"] == "kibana"
	})
	if existingKibanaEnv != nil {
		defaultUrl := (*existingKibanaEnv)["url"]
		fmt.Printf("URL [%s]: ", defaultUrl)
		em["url"] = readLine(reader)
		if em["url"] == "" {
			em["auth"] = (*existingKibanaEnv)["auth"]
			em["url"] = defaultUrl
		}
	} else {
		fmt.Print("URL: ")
		em["url"] = readLine(reader)
	}
	var kibanaClient *kibana.Client
	var indices []string
	var err error
	for {
		fmt.Println("Attempting to connect to Kibana on ", em["url"])
		kibanaClient = kibana.New(em["url"], em["auth"], "")
		indices, err = kibanaClient.ListIndices()
		if err != nil && err.Error() == "Authentication failed" {
			user, pass := credentials(reader)
			em["auth"] = fmt.Sprintf("Basic %s", b64Encode(fmt.Sprintf("%s:%s", user, pass)))
			continue
		} else if err != nil {
			fmt.Printf("Got error connecting to Kibana: %s\n", err)
			return em, err
		}
		break
	}
	fmt.Println("List of indices:")
	for _, index := range indices {
		fmt.Println("  ", index)
	}
	fmt.Print("Index: ")
	em["index"] = readLine(reader)
	return em, nil
}


<|start_filename|>pkg/backend/cloudwatch/client_test.go<|end_filename|>
package cloudwatch

import (
	"testing"

	"github.com/egnyte/ax/pkg/backend/common"
)

func TestParsing(t *testing.T) {
	msg := attemptParseJSON(`2017-09-27T09:01:01.245468966Z {"asctime": "2017-09-27 09:01:01,245", "created": 1506502861.2452097, "filename": "connectionpool.py", "funcName": "_make_request", "levelname": "DEBUG", "levelno": 10, "module": "connectionpool", "msecs": 245.2096939086914, "message": "http://localhost:None \"POST /v1.29/exec/1744fb9d8aa1ed1f94f729d4e0474251dfab9e0523385d42e77ea10acda53957/start HTTP/1.1\" 101 0", "name": "urllib3.connectionpool", "pathname": "/usr/local/lib/python3.6/site-packages/urllib3/connectionpool.py", "process": 5, "processName": "MainProcess", "relativeCreated": 2276.298999786377, "thread": 140018892404480, "threadName": "MainThread", "turbo_request_id": null, "user": null, "tid": 5, "source": "/usr/local/lib/python3.6/site-packages/urllib3/connectionpool.py:395", "client_id": null}
		`)
	if msg["filename"] != "connectionpool.py" {
		t.Errorf("Parsed: %+v", msg)
		t.Fail()
	}
	msg = attemptParseJSON(`{"asctime": "2017-09-27 09:01:01,245", "created": 1506502861.2452097, "filename": "connectionpool.py", "funcName": "_make_request", "levelname": "DEBUG", "levelno": 10, "module": "connectionpool", "msecs": 245.2096939086914, "message": "http://localhost:None \"POST /v1.29/exec/1744fb9d8aa1ed1f94f729d4e0474251dfab9e0523385d42e77ea10acda53957/start HTTP/1.1\" 101 0", "name": "urllib3.connectionpool", "pathname": "/usr/local/lib/python3.6/site-packages/urllib3/connectionpool.py", "process": 5, "processName": "MainProcess", "relativeCreated": 2276.298999786377, "thread": 140018892404480, "threadName": "MainThread", "turbo_request_id": null, "user": null, "tid": 5, "source": "/usr/local/lib/python3.6/site-packages/urllib3/connectionpool.py:395", "client_id": null}
		`)
	if msg["filename"] != "connectionpool.py" {
		t.Errorf("Parsed: %+v", msg)
		t.Fail()
	}
}

func TestFilterGenerator(t *testing.T) {
	output := queryToFilterPattern(common.Query{
		QueryString: "Test",
		EqualityFilters: []common.EqualityFilter{
			{
				FieldName: "name",
				Operator:  "=",
				Value:     "zef",
			},
		},
	})
	if output != `Test { ($.name = "zef") }` {
		t.Fatal(output)
	}

	output = queryToFilterPattern(common.Query{
		EqualityFilters: []common.EqualityFilter{
			{
				FieldName: "name",
				Operator:  "=",
				Value:     "zef",
			},
		},
	})
	if output != `{ ($.name = "zef") }` {
		t.Fatal(output)
	}

	output = queryToFilterPattern(common.Query{
		EqualityFilters: []common.EqualityFilter{
			{
				FieldName: "name",
				Operator:  "=",
				Value:     "zef",
			},
			{
				FieldName: "age",
				Operator:  "=",
				Value:     "34",
			},
		},
	})
	if output != `{ ($.name = "zef") && ($.age = "34") }` {
		t.Fatal(output)
	}
}


<|start_filename|>pkg/backend/stackdriver/client_test.go<|end_filename|>
package stackdriver

import (
	"fmt"
	"testing"

	"github.com/egnyte/ax/pkg/backend/common"
)

func TestAttributeDecoding(t *testing.T) {
	exampleJson := []byte(`{"fields":{"age":{"Kind":{"NumberValue":34}},"bool":{"Kind":{"BoolValue":true}},"list":{"Kind":{"ListValue":{"values":[{"Kind":{"NumberValue":1}},{"Kind":{"NumberValue":2}},{"Kind":{"NumberValue":3}},{"Kind":{"NumberValue":4}}]}}},"name":{"Kind":{"StringValue":"test"}},"obj":{"Kind":{"StructValue":{"fields":{"name":{"Kind":{"StringValue":"test"}}}}}},"slist":{"Kind":{"ListValue":{"values":[{"Kind":{"StringValue":"aap"}},{"Kind":{"StringValue":"noot"}},{"Kind":{"StringValue":"mies"}}]}}}}}`)
	m := payloadToAttributes(exampleJson)
	fmt.Printf("%+v\n", m)
	if m["age"] != int64(34) {
		t.Error("age")
	}
	if m["bool"] != true {
		t.Error("bool")
	}
	if len(m["list"].([]interface{})) != 4 {
		t.Error("list")
	}
	if len(m["slist"].([]interface{})) != 3 {
		t.Error("slist")
	}
	obj, ok := m["obj"].(map[string]interface{})
	if !ok {
		t.Error("obj")
	}

	if obj["name"] != "test" {
		t.Error("obj.name")
	}
}

func TestQueryToFilter(t *testing.T) {
	if queryToFilter(common.Query{}, "my-project", "my-log") != `logName = "projects/my-project/logs/my-log"` {
		t.Error("Empty search")
	}
	if queryToFilter(common.Query{QueryString: "My query"}, "my-project", "my-log") != `logName = "projects/my-project/logs/my-log" AND "My query"` {
		t.Error("Basic search filter")
	}
	if queryToFilter(common.Query{
		EqualityFilters: []common.EqualityFilter{
			{
				FieldName: "name",
				Operator:  "=",
				Value:     "pete",
			},
		}}, "my-project", "my-log") != `logName = "projects/my-project/logs/my-log" AND jsonPayload.name = "pete"` {
		t.Error("Where filter fail")
	}
}


<|start_filename|>pkg/alert/slack/slack.go<|end_filename|>
package slack

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	yaml "gopkg.in/yaml.v2"

	"github.com/egnyte/ax/pkg/alert"
	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/cache"
)

type SlackAlerter struct {
	name      string
	token     string
	channel   string
	username  string
	iconEmoji string
	seenCache *cache.Cache
}

func New(name, dataDir string, config map[string]string) *SlackAlerter {
	return &SlackAlerter{
		name:      name,
		token:     config["token"],
		channel:   config["channel"],
		username:  config["username"],
		iconEmoji: config["icon_emoji"],
		seenCache: cache.New(fmt.Sprintf("%s/alert-%s-seen.json", dataDir, name)),
	}
}

type slackResponse struct {
	Ok      bool
	Error   string
	Channel string
	Ts      string
}

func slackPostRequest(url string, form url.Values) (*slackResponse, error) {
	req, err := http.NewRequest("POST", url, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	decoder := json.NewDecoder(res.Body)
	var slackResponse slackResponse
	err = decoder.Decode(&slackResponse)
	if err != nil {
		return nil, err
	}
	if !slackResponse.Ok {
		return nil, errors.New(slackResponse.Error)
	}
	return &slackResponse, nil
}

func (alerter *SlackAlerter) SendAlert(lm common.LogMessage) error {
	contentHash := lm.ContentHash()
	buf, _ := yaml.Marshal(lm.Attributes)
	if alerter.seenCache.Contains(contentHash) {
		fmt.Println("Skipping", lm)
		/*
			ts := alerter.seenCache.GetString(contentHash)
			// TODO: Add reaction
			form := url.Values{}
			form.Add("token", alerter.token)
			form.Add("channel", alerter.channel)
			form.Add("ts", ts)
			form.Add("text", fmt.Sprintf("(:exclamation: %d) *[%s]* %s", occurrences, lm.Timestamp.Format(common.TimeFormat), buf))
			form.Add("username", alerter.username)
			fmt.Println("Skipped", lm)
			resp, err := slackPostRequest("https://slack.com/api/chat.update", form)
			fmt.Println("Response from slack", resp, err)
			return err
		*/
		return nil
	}
	form := url.Values{}
	form.Add("token", alerter.token)
	form.Add("icon_emoji", alerter.iconEmoji)
	form.Add("channel", alerter.channel)
	form.Add("text", fmt.Sprintf("*[%s]* %s", lm.Timestamp.Format(common.TimeFormat), buf))
	form.Add("username", alerter.username)
	resp, err := slackPostRequest("https://slack.com/api/chat.postMessage", form)
	if err != nil {
		return err
	}
	if !resp.Ok {
		return errors.New(resp.Error)
	}

	expire := time.Now().Add(time.Hour * 24 * 7)
	alerter.seenCache.Set(contentHash, resp.Ts, &expire)

	return alerter.seenCache.Flush()
}

var _ alert.Alerter = &SlackAlerter{}


<|start_filename|>pkg/config/colors.go<|end_filename|>
package config

import "github.com/fatih/color"

type colorDef struct {
	Fg        string
	Bg        string
	Bold      bool
	Faint     bool
	Italic    bool
	Underline bool
}

type ColorConfig struct {
	Timestamp      colorDef
	AttributeKey   colorDef
	AttributeValue colorDef
	Message        colorDef
}

var defaultColorConfig ColorConfig = ColorConfig{
	Timestamp: colorDef{
		Fg: "magenta",
	},
	Message: colorDef{
		Bold: true,
	},
	AttributeKey: colorDef{
		Fg: "cyan",
	},
}

func colorStringToAttribute(s string, fg bool) color.Attribute {
	fgOrBg := func(f color.Attribute, b color.Attribute) color.Attribute {
		if fg {
			return f
		} else {
			return b
		}
	}
	switch s {
	case "red":
		return fgOrBg(color.FgRed, color.BgRed)
	case "green":
		return fgOrBg(color.FgGreen, color.BgGreen)
	case "yellow":
		return fgOrBg(color.FgYellow, color.BgYellow)
	case "blue":
		return fgOrBg(color.FgBlue, color.BgBlue)
	case "magenta":
		return fgOrBg(color.FgMagenta, color.BgMagenta)
	case "cyan":
		return fgOrBg(color.FgCyan, color.BgCyan)
	case "white":
		return fgOrBg(color.FgWhite, color.BgWhite)
	}
	return color.Reset
}

func ColorToTermColor(colorDef colorDef) *color.Color {
	c := color.New()

	if colorDef.Fg != "" {
		c = c.Add(colorStringToAttribute(colorDef.Fg, true))
	}
	if colorDef.Bg != "" {
		c = c.Add(colorStringToAttribute(colorDef.Bg, false))
	}

	if colorDef.Bold {
		c = c.Add(color.Bold)
	}
	if colorDef.Faint {
		c = c.Add(color.Faint)
	}
	if colorDef.Italic {
		c = c.Add(color.Italic)
	}
	if colorDef.Underline {
		c = c.Add(color.Underline)
	}
	return c
}


<|start_filename|>pkg/backend/kibana/query_test.go<|end_filename|>
package kibana

import "testing"

func TestProject(t *testing.T) {
	/*myMap := project(map[string]interface{}{
		"name": "Hello",
		"age":  33,
		"docker": map[string]interface{}{
			"service": "bla",
		},
	}, []string{"age", "docker.service"})
	if !reflect.DeepEqual(myMap, map[string]interface{}{
		"age":            33,
		"docker.service": "bla",
	}) {
		t.Fatalf("Can't properly map: %+v", myMap)
	}
	*/
}


<|start_filename|>pkg/backend/kibana/util_test.go<|end_filename|>
package kibana

import "testing"

func TestSafeFilename(t *testing.T) {
	if safeFilename("turbo-*") != "turbo-_" {
		t.Fatal("Not a safe filename:", safeFilename("turbo-*"))
	}
}


<|start_filename|>Makefile<|end_filename|>
all:
	go install github.com/egnyte/ax/cmd/ax

test:
	go test ./cmd/... ./pkg/...

deps:
	dep ensure

release:
	rm -rf dist bin
	goreleaser


<|start_filename|>pkg/complete/complete.go<|end_filename|>
package complete

import (
	"fmt"
	"time"

	"log"

	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/cache"
	"github.com/egnyte/ax/pkg/config"
)

const cacheFilename = "attribute-cache.json"

func GatherCompletionInfo(rc config.RuntimeConfig, messages <-chan common.LogMessage) <-chan common.LogMessage {
	cache := cache.New(fmt.Sprintf("%s/%s", rc.DataDir, cacheFilename))
	completionsKey := fmt.Sprintf("completions:%s", rc.ActiveEnv)
	attrNames := make(map[string]bool)

	// This will be read back as a map[string]interface{} not a bool
	if existingAttributes, ok := cache.Get(completionsKey).(map[string]interface{}); ok {
		log.Println("Found existing attribute cache")
		for existingAttr := range existingAttributes {
			attrNames[existingAttr] = true
		}
	}
	changed := true
	stopFlushing := make(chan struct{})
	resultChan := make(chan common.LogMessage)
	go func() {
		for message := range messages {
			resultChan <- message
			for k := range message.Attributes {
				if !attrNames[k] {
					attrNames[k] = true
					changed = true
				}
			}
		}
		close(resultChan)
		stopFlushing <- struct{}{}
	}()
	go func() {
		// Flushes cache to disk every 5 seconds until completed
		for {
			shouldBreak := false
			select {
			case <-stopFlushing:
				shouldBreak = true
			case <-time.After(5 * time.Second):
			}
			if changed {
				log.Println("Flushing cache")
				cache.Set(completionsKey, attrNames, nil)
				err := cache.Flush()
				if err != nil {
					log.Println("Could not flush cache:", err)
				}
				changed = false
			}
			if shouldBreak {
				break
			}
		}
		log.Println("Stopped flush loop")
	}()
	return resultChan
}

func GetCompletions(rc config.RuntimeConfig) map[string]bool {
	cache := cache.New(fmt.Sprintf("%s/%s", rc.DataDir, cacheFilename))
	res := cache.Get(fmt.Sprintf("completions:%s", rc.ActiveEnv))
	result := make(map[string]bool)
	if attrNames, ok := res.(map[string]interface{}); ok {
		for attrName := range attrNames {
			result[attrName] = true
		}
	}
	return result
}


<|start_filename|>pkg/backend/docker/client.go<|end_filename|>
package docker

import (
	"context"
	"fmt"
	"log"
	"os/exec"
	"strings"

	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/backend/subprocess"
)

type DockerClient struct {
	containerPattern string
}

func GetRunningContainers(pattern string) []string {
	flags := []string{"ps", "--format", "{{.Names}}"}
	if pattern != "" {
		flags = append(flags, "--filter", fmt.Sprintf("name=%s", pattern))
	}
	allContainers, err := exec.Command("docker", flags...).Output()
	if err != nil {
		log.Printf("Retrieving all containers failed: %v\n", err)
		return []string{}
	}
	return strings.Split(strings.TrimSpace(string(allContainers)), "\n")
}

func DockerHintAction() []string {
	return GetRunningContainers("")
}

func (client *DockerClient) ImplementsAdvancedFilters() bool {
	return true
}

func (client *DockerClient) Query(ctx context.Context, query common.Query) <-chan common.LogMessage {
	resultChan := make(chan common.LogMessage)
	runningCommands := 0
	for _, containerName := range GetRunningContainers(client.containerPattern) {
		command := []string{"docker", "logs", "--tail", fmt.Sprintf("%d", query.MaxResults)}
		if query.Follow {
			command = append(command, "-f")
		}
		command = append(command, containerName)
		client := subprocess.New(command)
		runningCommands++
		go func() {
			for message := range client.Query(ctx, query) {
				message.Attributes["@container"] = containerName
				resultChan <- message
			}
			runningCommands--
			if runningCommands == 0 {
				close(resultChan)
			}
		}()
	}
	return resultChan
}

func New(containerPattern string) *DockerClient {
	return &DockerClient{containerPattern}
}

var _ common.Client = &DockerClient{}


<|start_filename|>cmd/ax/alert_test.go<|end_filename|>
package main

import (
	"testing"
)

func TestSlack(t *testing.T) {
	//alertSlack(os.Getenv("SLACK_TOKEN"), "#zeftest", "Sup from Go")
}


<|start_filename|>pkg/backend/common/duplicate_filter_test.go<|end_filename|>
package common

import (
	"testing"
)

func TestDedup(t *testing.T) {
	inChannel := make(chan LogMessage, 10)
	noDupChannel := Dedup(inChannel)
	inChannel <- LogMessage{
		ID: "1",
	}
	inChannel <- LogMessage{
		ID: "2",
	}
	inChannel <- LogMessage{
		ID: "3",
	}
	if msg := <-noDupChannel; msg.ID != "1" {
		t.Fatal("Didn't get first message")
	}
	if msg := <-noDupChannel; msg.ID != "2" {
		t.Fatal("Didn't get second message")
	}
	if msg := <-noDupChannel; msg.ID != "3" {
		t.Fatal("Didn't get third message")
	}
	inChannel <- LogMessage{
		ID: "3",
	}
	inChannel <- LogMessage{
		ID: "2",
	}
	inChannel <- LogMessage{
		ID: "4",
	}
	if msg := <-noDupChannel; msg.ID != "4" {
		t.Fatal("Should have received 4")
	}
}


<|start_filename|>pkg/backend/stackdriver/client.go<|end_filename|>
package stackdriver

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"time"

	"google.golang.org/api/option"

	"cloud.google.com/go/logging"
	"cloud.google.com/go/logging/logadmin"
	"github.com/egnyte/ax/pkg/backend/common"
	"google.golang.org/api/iterator"
)

const QueryLogTimeout = 20 * time.Second

type StackdriverClient struct {
	stackdriverClient *logadmin.Client
	projectName       string
	logName           string
}

// This is some crazy-ass structure in which the stackdriver APIs
// return its JSON values that we have to decode
type payloadValue struct {
	Fields map[string]payloadEntry
}

type payloadEntry struct {
	Kind struct {
		StringValue *string
		NumberValue *int64
		BoolValue   *bool
		ListValue   *struct {
			Values []payloadEntry
		}
		StructValue *payloadValue
	}
}

func payloadValueToJSONValue(plVal payloadValue) map[string]interface{} {
	m := make(map[string]interface{})
	for k, v := range plVal.Fields {
		m[k] = payloadEntryToJSONValue(v)
	}
	return m
}

func payloadEntryToJSONValue(plEntry payloadEntry) interface{} {
	kind := plEntry.Kind
	if kind.StringValue != nil {
		return *plEntry.Kind.StringValue
	} else if kind.NumberValue != nil {
		return *kind.NumberValue
	} else if kind.BoolValue != nil {
		return *kind.BoolValue
	} else if kind.ListValue != nil {
		list := make([]interface{}, len((*kind.ListValue).Values))
		for idx, val := range (*kind.ListValue).Values {
			list[idx] = payloadEntryToJSONValue(val)
		}
		return list
	} else if kind.StructValue != nil {
		return payloadValueToJSONValue(*kind.StructValue)
	} else {
		return nil
	}
}

func payloadToAttributes(buf []byte) map[string]interface{} {
	var plValue payloadValue
	if err := json.Unmarshal(buf, &plValue); err != nil {
		fmt.Printf("Could not unmarshall value: %s", string(buf))
		return nil
	}
	return payloadValueToJSONValue(plValue)
}

func entryToLogMessage(entry *logging.Entry) common.LogMessage {
	message := common.NewLogMessage()
	message.Timestamp = entry.Timestamp
	message.ID = entry.InsertID
	switch v := entry.Payload.(type) {
	case string:
		message.Attributes["message"] = v
	case map[string]interface{}:
		message.Attributes = v
	default:
		buf, err := json.Marshal(entry.Payload)
		if err != nil {
			fmt.Printf("Could not marshall value: %v of type %v", entry.Payload, reflect.TypeOf(entry))
			break
		}
		message.Attributes = payloadToAttributes(buf)
	}
	return message
}

func queryToFilter(query common.Query, projectName string, logName string) string {
	pieces := []string{fmt.Sprintf(`logName = "projects/%s/logs/%s"`, projectName, logName)}
	if query.QueryString != "" {
		pieces = append(pieces, fmt.Sprintf(`"%s"`, query.QueryString))
	}
	for _, filter := range query.EqualityFilters {
		pieces = append(pieces, fmt.Sprintf(`jsonPayload.%s %s "%s"`, filter.FieldName, filter.Operator, filter.Value))
	}
	if query.After != nil {
		pieces = append(pieces, fmt.Sprintf(`timestamp > "%s"`, (*query.After).Format(time.RFC3339)))
	}
	if query.Before != nil {
		pieces = append(pieces, fmt.Sprintf(`timestamp < "%s"`, (*query.Before).Format(time.RFC3339)))
	}
	return strings.Join(pieces, " AND ")
}

func (client *StackdriverClient) ImplementsAdvancedFilters() bool {
	return false
}

func (client *StackdriverClient) readLogBatch(ctx context.Context, query common.Query) ([]common.LogMessage, error) {
	ctx, cancel := context.WithTimeout(ctx, QueryLogTimeout)
	defer cancel()
	it := client.stackdriverClient.Entries(ctx, logadmin.Filter(queryToFilter(query, client.projectName, client.logName)))
	messages := make([]common.LogMessage, 0, 20)
	entry, err := it.Next()
	// Somehow, if no results can be found, it.Next() just runs forever, hence the adding a timeout to the context
	if ctx.Err() == context.DeadlineExceeded {
		return messages, ctx.Err()
	}
	if err != nil && err != iterator.Done {
		return nil, err
	}
	resultCounter := 1
	for err != iterator.Done && resultCounter <= query.MaxResults {
		msg := entryToLogMessage(entry)
		msg.Attributes = common.Project(msg.Attributes, query.SelectFields)
		messages = append(messages, msg)
		entry, err = it.Next()
		resultCounter++
	}
	return messages, nil
}

func (client *StackdriverClient) Query(ctx context.Context, query common.Query) <-chan common.LogMessage {
	if query.Follow {
		return common.ReQueryFollow(ctx, func() ([]common.LogMessage, error) {
			return client.readLogBatch(ctx, query)
		})
	}
	resultChan := make(chan common.LogMessage)

	go func() {
		messages, err := client.readLogBatch(ctx, query)
		if err != nil {
			fmt.Printf("Error while fetching logs: %s\n", err)
			close(resultChan)
			return
		}
		for _, message := range messages {
			resultChan <- message
		}
		close(resultChan)
	}()

	return resultChan
}

func New(credentialsFile, projectName, logName string) *StackdriverClient {
	client, err := logadmin.NewClient(context.Background(), projectName, option.WithCredentialsFile(credentialsFile))
	if err != nil {
		fmt.Printf("Error creating stack driver client: %v\n", err)
		return nil
	}
	return &StackdriverClient{
		stackdriverClient: client,
		projectName:       projectName,
		logName:           logName,
	}
}

func (client *StackdriverClient) ListLogs() ([]string, error) {
	logNames := make([]string, 0, 10)
	it := client.stackdriverClient.Logs(context.Background())
	s, err := it.Next()
	if err != nil && err != iterator.Done {
		return nil, err
	}
	for err != iterator.Done {
		logNames = append(logNames, s)
		s, err = it.Next()
		if err != nil && err != iterator.Done {
			return nil, err
		}
	}
	return logNames, nil
}

var _ common.Client = &StackdriverClient{}


<|start_filename|>pkg/backend/subprocess/client.go<|end_filename|>
package subprocess

import (
	"context"
	"fmt"
	"log"
	"os/exec"

	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/backend/stream"
)

type SubprocessClient struct {
	command []string
}

func (client *SubprocessClient) ImplementsAdvancedFilters() bool {
	return true
}

func (client *SubprocessClient) Query(ctx context.Context, query common.Query) <-chan common.LogMessage {
	resultChan := make(chan common.LogMessage)
	cmd := exec.Command(client.command[0], client.command[1:]...)
	stdOut, err := cmd.StdoutPipe()
	if err != nil {
		log.Printf("Could not get stdout pipe: %v", err)
		close(resultChan)
		return resultChan
	}
	stdOutStream := stream.New(stdOut)
	stdErr, err := cmd.StderrPipe()
	if err != nil {
		log.Printf("Could not get stderr pipe: %v", err)
		close(resultChan)
		return resultChan
	}
	stdErrStream := stream.New(stdErr)
	if err := cmd.Start(); err != nil {
		fmt.Printf("Could not start process: %s because: %v\n", client.command[0], err)
		close(resultChan)
		return resultChan
	}
	go func() {
		stdOutQuery := stdOutStream.Query(ctx, query)
		stdErrQuery := stdErrStream.Query(ctx, query)
		stdOutOpen := true
		stdErrOpen := true
		for stdOutOpen || stdErrOpen {
			select {
			case message, ok := <-stdOutQuery:
				if !ok {
					stdOutOpen = false
					continue
				}
				resultChan <- message
			case message, ok := <-stdErrQuery:
				if !ok {
					stdErrOpen = false
					continue
				}
				resultChan <- message
			case <-ctx.Done():
				close(resultChan)
				cmd.Process.Kill() // Ignoring error, not sure if that's ok
				// Returning to avoid the Wait()
				return
			}

		}
		close(resultChan)
		if err := cmd.Wait(); err != nil {
			fmt.Printf("Process exited with error: %v\n", err)
		}
	}()
	return resultChan
}

func New(command []string) *SubprocessClient {
	return &SubprocessClient{command}
}

var _ common.Client = &SubprocessClient{}


<|start_filename|>pkg/config/stackdriver.go<|end_filename|>
package config

import (
	"bufio"
	"fmt"

	"github.com/egnyte/ax/pkg/backend/stackdriver"
)

func stackdriverConfig(reader *bufio.Reader, existingConfig Config) (EnvMap, error) {
	em := EnvMap{
		"backend": "stackdriver",
	}
	existingSDEnv := findFirstEnvWhere(existingConfig.Environments, func(em EnvMap) bool {
		return em["backend"] == "stackdriver"
	})
	if existingSDEnv != nil {
		credentialsPath := (*existingSDEnv)["credentials"]
		fmt.Printf("Path to credentials file (JSON) [%s]: ", credentialsPath)
		em["credentials"] = readLine(reader)
		if em["credentials"] == "" {
			em["credentials"] = (*existingSDEnv)["credentials"]
		}
	} else {
		fmt.Print("Path to credentials file (JSON): ")
		em["credentials"] = readLine(reader)
	}
	fmt.Print("GCP Project name: ")
	em["project"] = readLine(reader)
	var sdClient *stackdriver.StackdriverClient
	var logs []string
	var err error
	fmt.Println("Attempting to connect to Stackdriver")
	sdClient = stackdriver.New(em["credentials"], em["project"], "")
	logs, err = sdClient.ListLogs()
	if err != nil {
		fmt.Printf("Got error connecting to Stackdriver: %s\n", err)
		return em, err
	}
	fmt.Println("List of logs:")
	for _, log := range logs {
		fmt.Println("  ", log)
	}
	fmt.Print("Log: ")
	em["log"] = readLine(reader)
	return em, nil
}


<|start_filename|>pkg/cache/cache_test.go<|end_filename|>
package cache

import (
	"testing"
	"time"
)

func TestCacheBehavior(t *testing.T) {
	cache := New("test.cache")
	if cache.Contains("test-key") {
		t.Error("Key should not be present")
	}
	cache.Set("test-key", "value", nil)
	if !cache.Contains("test-key") {
		t.Error("Key should be present")
	}
	if cache.Get("test-key") != "value" {
		t.Errorf("Key has wrong value: %v", cache.Get("test-key"))
	}
	if cache.Flush() != nil {
		t.Errorf("Flush failed")
	}
	cache.Unset("test-key")
	if cache.Contains("test-key") {
		t.Error("Key should not be present")
	}
	cache.Remove()
}

func TestCacheExpiry(t *testing.T) {
	cache := New("test.cache")
	expireDate := time.Now().Add(time.Second)
	cache.Set("test-key", "hello", &expireDate)
	if !cache.Contains("test-key") {
		t.Error("Key not present")
	}
	time.Sleep(time.Second * 2)
	if cache.Contains("test-key") {
		t.Error("Key should have expired")
	}
	// cache.Remove()
}

func TestCachePersistence(t *testing.T) {

}


<|start_filename|>cmd/ax/upgrade.go<|end_filename|>
package main

import (
	"archive/tar"
	"bufio"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/google/go-github/github"
	"github.com/kardianos/osext"
)

var (
	gitOrganisationName = "egnyte"
	gitRepositoryName   = "ax"
)

func getLatestReleaseData() *github.RepositoryRelease {
	client := github.NewClient(nil)

	release, _, err := client.Repositories.GetLatestRelease(context.Background(), gitOrganisationName, gitRepositoryName)
	if err != nil {
		fmt.Printf("Couldn't fetch release information: %v", err)
		os.Exit(1)
	}

	return release
}

func getLatestAssetLink() (string, error) {
	release := getLatestReleaseData()

	for _, element := range release.Assets {
		assetLink := element.GetBrowserDownloadURL()
		if strings.Contains(assetLink, runtime.GOOS) && strings.Contains(assetLink, runtime.GOARCH) {
			return assetLink, nil
		}
	}

	return "", errors.New("assset link couldn't be found")
}

func getLatestReleaseTag() string {
	release := getLatestReleaseData()
	return release.GetTagName()
}

func downloadFile(filePath string, url string) error {

	out, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer out.Close()

	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if _, err = io.Copy(out, resp.Body); err != nil {
		return err
	}

	return nil
}

func upgradeVersion() error {
	latestVersion := getLatestReleaseTag()

	if version == latestVersion {
		fmt.Println("Latest version is already installed.")
		return nil
	}

	fmt.Printf("New version detected. Current version: %s. Latest version: %s\n", version, latestVersion)
	inputScanner := bufio.NewScanner(os.Stdin)
	fmt.Print("Do you want to upgrade Ax to the latest version? (yes/no): ")
	inputScanner.Scan()

	if text := inputScanner.Text(); strings.ToLower(text) != "yes" {
		fmt.Println("Aborting the upgrade.")
		os.Exit(1)
	}

	latestAssetURL, err := getLatestAssetLink()
	if err != nil {
		return err
	}

	extractionPath, err := ioutil.TempDir("/tmp", "egnyte_ax")
	if err != nil {
		return err
	}
	downloadPath := fmt.Sprintf("%s/ax.tar.gz", extractionPath)

	if _, err := os.Stat(extractionPath); os.IsNotExist(err) {
		os.Mkdir(extractionPath, 0755)
	}

	fmt.Println("Ax upgrade in progress...")

	if err := downloadFile(downloadPath, latestAssetURL); err != nil {
		return err
	}

	file, err := os.Open(downloadPath)
	if err != nil {
		return err
	}

	// extract the tar.gz file
	if err := extractTar(extractionPath, file); err != nil {
		return err
	}

	// switch old binary with a new one
	currentBinaryPath, _ := osext.Executable()
	if err := os.Rename(extractionPath+"/ax", currentBinaryPath); err != nil {
		return err
	}
	defer os.RemoveAll(extractionPath)

	return nil
}

func extractTar(dst string, r io.Reader) error {

	gzr, err := gzip.NewReader(r)
	defer gzr.Close()
	if err != nil {
		return err
	}

	tr := tar.NewReader(gzr)

	for {
		header, err := tr.Next()

		switch {

		// if no more files are found return
		case err == io.EOF:
			return nil

		// return any other error
		case err != nil:
			return err

		// if the header is nil, just skip it (not sure how this happens)
		case header == nil:
			continue
		}

		// the target location where the dir/file should be created
		target := filepath.Join(dst, header.Name)

		// check the file type
		switch header.Typeflag {

		// if its a dir and it doesn't exist create it
		case tar.TypeDir:
			if _, err := os.Stat(target); err != nil {
				if err := os.MkdirAll(target, 0755); err != nil {
					return err
				}
			}

		// if it's a file create it
		case tar.TypeReg:
			f, err := os.OpenFile(target, os.O_CREATE|os.O_RDWR, os.FileMode(header.Mode))
			if err != nil {
				return err
			}
			defer f.Close()

			// copy over contents
			if _, err := io.Copy(f, tr); err != nil {
				return err
			}
		}
	}
}


<|start_filename|>pkg/config/config.go<|end_filename|>
package config

import (
	"bufio"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/imdario/mergo"
	"github.com/zefhemel/kingpin"
	yaml "gopkg.in/yaml.v2"

	"github.com/egnyte/ax/pkg/backend/common"
	"github.com/egnyte/ax/pkg/backend/docker"
	"github.com/olekukonko/tablewriter"
)

var dataDir string

type EnvMap map[string]string

type Config struct {
	DefaultEnv   string            `yaml:"default"`
	Colors       ColorConfig       `yaml:"colors"`
	Environments map[string]EnvMap `yaml:"env"`
	Alerts       []AlertConfig     `yaml:"alerts"`
}

type AlertConfig struct {
	Env      string                `yaml:"env"`
	Name     string                `yaml:"name"`
	Selector common.QuerySelectors `yaml:"selector"`
	Service  AlertServiceConfig    `yaml:"service"`
}

type AlertServiceConfig map[string]string

type RuntimeConfig struct {
	ActiveEnv string
	DataDir   string
	Env       EnvMap
	Config    Config
}

var (
	activeEnv      = kingpin.Flag("env", "Environment to connect to").Short('e').HintAction(envHintAction).String()
	dockerFlag     = kingpin.Flag("docker", "Query docker container logs").HintAction(docker.DockerHintAction).String()
	envCommand     = kingpin.Command("env", "Environment management commands")
	envInitCommand = envCommand.Command("add", "Add an environment")
	envEditCommand = envCommand.Command("edit", "Edit your environment configuration file in a text editor")
	envListCommand = envCommand.Command("list", "List all environments").Default()
)

func NewConfig() Config {
	return Config{
		Environments: make(map[string]EnvMap),
		Alerts:       make([]AlertConfig, 0),
	}
}

func LoadConfig() Config {
	config := NewConfig()
	buf, err := ioutil.ReadFile(configPathName())
	if err != nil {
		return config
	}
	err = yaml.UnmarshalStrict(buf, &config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Could not unmarshall config: %s", err)
		return config
	}
	if config.Environments == nil {
		config.Environments = make(map[string]EnvMap)
	}
	if config.Alerts == nil {
		config.Alerts = make([]AlertConfig, 0)
	}
	if err := mergo.Merge(&config.Colors, defaultColorConfig); err != nil {
		panic("Could not set default colors")
	}
	return config
}

func configPathName() string {
	return fmt.Sprintf("%s/ax.yaml", dataDir)
}

func BuildConfig() RuntimeConfig {
	config := LoadConfig()
	rc := RuntimeConfig{
		DataDir: dataDir,
		Env:     make(EnvMap),
		Config:  config,
	}
	var ok bool
	if config.DefaultEnv != "" {
		rc.Env, ok = config.Environments[config.DefaultEnv]
		if !ok {
			fmt.Println("Undefined active environment:", config.DefaultEnv)
			os.Exit(1)
		}
	}
	if *activeEnv != "" {
		rc.Env, ok = config.Environments[*activeEnv]
		rc.ActiveEnv = *activeEnv
		if !ok {
			fmt.Println("Undefined active environment:", *activeEnv)
			os.Exit(1)
		}
	}
	if *dockerFlag != "" {
		rc.ActiveEnv = fmt.Sprintf("docker.%s", *dockerFlag)
		rc.Env["backend"] = "docker"
		rc.Env["pattern"] = *dockerFlag
	}
	return rc
}

func credentials(reader *bufio.Reader) (string, string) {
	fmt.Print("Username: ")
	username, _ := reader.ReadString('\n')

	fmt.Print("Password: ")
	bytePassword, _ := terminal.ReadPassword(int(syscall.Stdin))
	fmt.Println()
	password := string(bytePassword)

	return strings.TrimSpace(username), strings.TrimSpace(password)
}

func readLine(reader *bufio.Reader) string {
	val, _ := reader.ReadString('\n')
	return strings.TrimSpace(val)
}

func b64Encode(s string) string {
	return base64.StdEncoding.EncodeToString([]byte(s))
}

func findFirstEnvWhere(environments map[string]EnvMap, whereFunc func(EnvMap) bool) *EnvMap {
	for _, m := range environments {
		if whereFunc(m) {
			return &m
		}
	}
	return nil
}

func SaveConfig(config Config) {
	f, err := os.Create(fmt.Sprintf("%s/ax.yaml", dataDir))
	if err != nil {
		fmt.Println("Couldn't open ax.yaml for writing", err)
		return
	}
	defer f.Close()
	buf, err := yaml.Marshal(&config)
	if err != nil {
		fmt.Println("Couldn't write to ax.yaml", err)
	}
	_, err = f.Write(buf)
	if err != nil {
		fmt.Println("Couldn't write to ax.yaml", err)
	}
}

func AddEnv() {
	config := LoadConfig()
	reader := bufio.NewReader(os.Stdin)
	fmt.Print("Name for new environment [default]: ")
	name := readLine(reader)
	if name == "" {
		name = "default"
	}
	fmt.Print("Choose a backend (kibana,cloudwatch,stackdriver) [kibana]: ")
	backend := readLine(reader)
	if backend == "" {
		backend = "kibana"
	}
	var em EnvMap
	var err error
	switch backend {
	case "kibana":
		em, err = kibanaConfig(reader, config)
		if err != nil {
			return
		}
	case "cloudwatch":
		em, err = cloudwatchConfig(reader, config)
		if err != nil {
			return
		}
	case "stackdriver":
		em, err = stackdriverConfig(reader, config)
		if err != nil {
			return
		}
	default:
		fmt.Println("Unsupported backend")
		return
	}
	if config.DefaultEnv == "" {
		config.DefaultEnv = name
	}
	config.Environments[name] = em
	SaveConfig(config)
}

func envHintAction() []string {
	config := LoadConfig()
	results := make([]string, 0, len(config.Environments))
	for k := range config.Environments {
		results = append(results, k)
	}
	return results
}
func ListEnvs() {
	config := LoadConfig()
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"D", "Name", "Backend", "Index"})
	for k, v := range config.Environments {
		def := ""
		if config.DefaultEnv == k {
			def = "*"
		}
		table.Append([]string{def, k, v["backend"], v["index"]})
	}
	table.Render() // Send output
}

func EditConfig() {
	editor := os.Getenv("EDITOR")
	if editor == "" {
		editor = "nano"
	}
	cmd := exec.Command(editor, configPathName())
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		fmt.Println("Error starting editor", err)
		return
	}
	if err := cmd.Wait(); err != nil {
		fmt.Println("Error waiting for editor", err)
		return
	}
}

func init() {
	dataDir = fmt.Sprintf("%s/.config/ax", os.Getenv("HOME"))
	err := os.MkdirAll(dataDir, 0700)
	if err != nil {
		fmt.Println("Could not create", dataDir)
		os.Exit(1)
	}

	// Set up logging
	logPath := fmt.Sprintf("%s/ax.log", dataDir)
	f, err := os.Create(logPath)
	if err != nil {
		fmt.Printf("Could not use %s for logging, logging to stdout instead\n", logPath)
		return // Skips log.SetOutput, defaults to stdout
	}
	log.SetOutput(f)
}


<|start_filename|>pkg/backend/kibana/query.go<|end_filename|>
package kibana

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sort"

	"github.com/egnyte/ax/pkg/backend/common"

	"os"

	"time"
)

type JsonObject map[string]interface{}
type JsonList []interface{}

type QueryResult struct {
	Responses []struct {
		Hits struct {
			Hits []Hit `json:"hits"`
		} `json:"hits"`
	} `json:"responses"`
}

type Hit struct {
	ID     string     `json:"_id"`
	Source JsonObject `json:"_source"`
}

type hitsByAscDate []Hit

func (a hitsByAscDate) Len() int      { return len(a) }
func (a hitsByAscDate) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a hitsByAscDate) Less(i, j int) bool {
	return a[i].Source["@timestamp"].(string) < a[j].Source["@timestamp"].(string)
}

func (client *Client) queryMessages(ctx context.Context, subIndex string, query common.Query) ([]Hit, error) {
	queryString := fmt.Sprintf("\"%s\"", query.QueryString) // TODO: Handle quotes properly
	if query.QueryString == "" {
		queryString = "*"
	}
	mustFilters := JsonList{
		JsonObject{
			"query_string": JsonObject{
				"analyze_wildcard": true,
				"query":            queryString,
			},
		},
	}

	if query.After != nil || query.Before != nil {
		rangeObj := JsonObject{
			"range": JsonObject{
				"@timestamp": JsonObject{
					"format": "epoch_millis",
				},
			},
		}
		if query.After != nil {
			rangeObj["range"].(JsonObject)["@timestamp"].(JsonObject)["gt"] = unixMillis(*query.After)
		}
		if query.Before != nil {
			rangeObj["range"].(JsonObject)["@timestamp"].(JsonObject)["lt"] = unixMillis(*query.Before)
		}
		mustFilters = append(mustFilters, rangeObj)
	}
	mustNotFilters := JsonList{}
	for _, filter := range query.EqualityFilters {
		m := JsonObject{}
		switch filter.Operator {
		case "=":
			m[filter.FieldName] = JsonObject{
				"query": filter.Value,
				"type":  "phrase",
			}
			mustFilters = append(mustFilters, JsonObject{
				"match": m,
			})
		case "!=":
			m[filter.FieldName] = JsonObject{
				"query": filter.Value,
				"type":  "phrase",
			}
			mustNotFilters = append(mustNotFilters, JsonObject{
				"match": m,
			})
		}
	}
	body, err := createMultiSearch(
		JsonObject{
			"index":              JsonList{subIndex},
			"ignore_unavailable": true,
		},
		JsonObject{
			"size": query.MaxResults,
			"sort": JsonList{
				JsonObject{
					"@timestamp": JsonObject{
						"order":         "desc",
						"unmapped_type": "boolean",
					},
				},
			},
			"query": JsonObject{
				"bool": JsonObject{
					"must":     mustFilters,
					"must_not": mustNotFilters,
				},
			},
		})
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", fmt.Sprintf("%s/elasticsearch/_msearch", client.URL), body)
	req = req.WithContext(ctx)
	if err != nil {
		return nil, err
	}
	client.addHeaders(req)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	var data QueryResult
	err = decoder.Decode(&data)
	if err != nil {
		return nil, err
	}
	hits := data.Responses[0].Hits.Hits
	sort.Sort(hitsByAscDate(hits))
	return hits, nil
}

// Implements "follow" mode for Kibana.
// Effectively this repeats the query every 5s and skips messages already seen
// Previously this was implemented by only requesting messages with a timestamp
// after the last seen one, but because sometimes logs arrive out of order, this
// resulted in skipping logs.
func (client *Client) queryFollow(ctx context.Context, q common.Query) <-chan common.LogMessage {
	return common.ReQueryFollow(ctx, func() ([]common.LogMessage, error) {
		return client.querySubIndex(ctx, client.Index, q)
	})
}

func (client *Client) Query(ctx context.Context, q common.Query) <-chan common.LogMessage {
	resultChan := make(chan common.LogMessage)
	if q.Before == nil {
		before := time.Now().Add(12 * time.Hour)
		q.Before = &before // Limit sanity
	}
	if q.Follow {
		return client.queryFollow(ctx, q)
	}
	go func() {
		printedResultsCount := 0
		fmt.Fprintf(os.Stderr, "Querying index %s\n", client.Index)
		allMessages, err := client.querySubIndex(ctx, client.Index, q)
		// Check if the context wasn't canceled
		select {
		case <-ctx.Done():
			close(resultChan)
			return
		default:
		}
		if err != nil {
			fmt.Fprintf(os.Stderr, "Could not connect to Kibana: %v", err)
			os.Exit(2)
		}
		for _, message := range allMessages {
			resultChan <- message
			printedResultsCount++
			if printedResultsCount >= q.MaxResults {
				break
			}
		}
		close(resultChan)
	}()

	return resultChan
}

func (client *Client) querySubIndex(ctx context.Context, subIndex string, q common.Query) ([]common.LogMessage, error) {
	hits, err := client.queryMessages(ctx, subIndex, q)
	if err != nil {
		return nil, err
	}

	allMessages := make([]common.LogMessage, 0, 200)
	for _, hit := range hits {
		attributes := hit.Source
		ts, err := time.Parse(time.RFC3339, attributes["@timestamp"].(string))
		if err != nil {
			return nil, err
		}
		delete(attributes, "@timestamp")
		message := common.FlattenLogMessage(common.LogMessage{
			ID:         hit.ID,
			Timestamp:  ts,
			Attributes: attributes,
		})
		message.Attributes = common.Project(message.Attributes, q.SelectFields)
		allMessages = append(allMessages, message)
	}
	return allMessages, nil
}


<|start_filename|>pkg/backend/common/duplicate_filter.go<|end_filename|>
package common

// Skips duplicate messages (based on .ID)
func Dedup(messageChan chan LogMessage) chan LogMessage {
	resultChan := make(chan LogMessage)
	idCache := make(map[string]bool)
	go func() {
		for message := range messageChan {
			if !idCache[message.ID] {
				resultChan <- message
				idCache[message.ID] = true
			}
		}
		close(resultChan)
	}()
	return resultChan
}
