package util import ( "context" log "github.com/sirupsen/logrus" "github.com/spf13/viper" "io/ioutil" "iwarma.ru/console/correlator/config" "iwarma.ru/console/correlator/es" "net/http" "testing" "time" ) func PrepareElastic() { viper.Set(config.ElasticUrl, "http://elasticsearch:9200") viper.Set(config.ElasticRetryCount, 1) viper.Set(config.ElasticUsername, "elastic") viper.Set(config.ElasticPassword, "changeme") viper.Set(config.Verbose, true) viper.Set(config.ElasticAggregatedIndexName, "test-aggregated") viper.Set(config.ElasticNormalizedIndexName, "test-normalized") viper.Set(config.AggregatorIterationDuration, time.Second*2) viper.Set(config.Threads, 10) eventsFields := []string{ "event_severity", "event_protocol", "message", "device_vendor", "device_product", "device_action", "device_version", "device_timezone", "sign_id", "sign_category", "sign_subcategory", "application", "source_ip", "source_host", "source_port", "source_mac", "source_timezone", "source_software", "source_action", "destination_ip", "destination_mac", "destination_timezone", "destination_software", "destination_action", "destination_host", "destination_port", "destination_user", "cs1", "cs1Label", "cs2", "cs2Label", "object_type", } viper.Set(config.AggregatedFields, eventsFields) } // TestUrl One server endpoint description type TestUrl struct { Url string // Url to handle Response []byte // Response to send ResponseHeaders map[string]string // Response headers to send ResponseStatus int // Status code for response. Default is 200 ReceivedMessage []byte // Received message from client Error error // Error from handler function CallCount int // Counter of requests to this endpoint } // TestServer Test HTTP server type TestServer struct { ServerAddress string // Address for server Urls []*TestUrl // Url endpoints Server *http.Server // Http server object ServerMux *http.ServeMux // Http multiplexer for server } // GetRequestAsString Get client request as string func (url TestUrl) GetRequestAsString() string { return string(url.ReceivedMessage) } // Start test server func (server *TestServer) Start() error { return server.Server.ListenAndServe() } // Stop test server func (server *TestServer) Stop() error { return server.Server.Shutdown(context.Background()) } // AddUrl Add url to server func (server *TestServer) AddUrl(url string, response []byte, responseHeaders map[string]string) *TestUrl { result := new(TestUrl) result.Url = url result.Response = response result.ResponseHeaders = responseHeaders result.ResponseStatus = http.StatusOK handler := func(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() result.ReceivedMessage, result.Error = ioutil.ReadAll(r.Body) result.CallCount = result.CallCount + 1 for key, value := range result.ResponseHeaders { w.Header().Set(key, value) } w.Write(result.Response) } server.ServerMux.HandleFunc(result.Url, handler) server.Urls = append(server.Urls, result) return result } func (server *TestServer) AddUrlWithStatus(url string, status int) *TestUrl { result := new(TestUrl) result.Url = url result.ResponseStatus = status handler := func(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() result.ReceivedMessage, result.Error = ioutil.ReadAll(r.Body) result.CallCount = result.CallCount + 1 w.WriteHeader(result.ResponseStatus) } server.ServerMux.HandleFunc(result.Url, handler) server.Urls = append(server.Urls, result) return result } // NewTestServer Create new test server func NewTestServer(address string) *TestServer { result := new(TestServer) result.ServerAddress = address result.Urls = make([]*TestUrl, 0) result.ServerMux = http.NewServeMux() result.Server = &http.Server{Addr: address, Handler: result.ServerMux} return result } func SetupTest(t *testing.T) { PrepareElastic() } func ClearIndex(el *es.Elastic, index ...string) error { for _, cur := range index { exists, err := el.CheckIndex(cur) if err != nil { log.Errorf("%+v", err) return err } if exists { err = el.DeleteIndex(cur) if err != nil { return err } } } return nil } func TearDownTest(t *testing.T) { client, err := es.NewClient() if err != nil { t.Error(err) return } del, err := client.DeleteIndex("*").Do(context.Background()) if err != nil { t.Errorf("%v", err) } if !del.Acknowledged { t.Errorf("Got bad response. Indexs, wasn't deleted") } }