Feature/add websocket server

cmd/api.go

@@ -0,0 +1,214 @@
+/*
+Copyright © 2021 Bisohns
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"os"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/bisohns/saido/config"
+	"github.com/bisohns/saido/driver"
+	"github.com/bisohns/saido/inspector"
+	"github.com/gorilla/handlers"
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+)
+
+const (
+	socketBufferSize  = 1042
+	messageBufferSize = 256
+)
+
+var (
+	port     string
+	server   = http.NewServeMux()
+	upgrader = &websocket.Upgrader{
+		ReadBufferSize:  socketBufferSize,
+		WriteBufferSize: socketBufferSize,
+		CheckOrigin: func(r *http.Request) bool {
+			return true
+		}}
+)
+
+type FullMessage struct {
+	Error   bool
+	Message interface{}
+}
+
+type Message struct {
+	Host string
+	Name string
+	Data interface{}
+}
+
+type Client struct {
+	Socket *websocket.Conn
+	Send   chan *FullMessage
+}
+
+// Write to websocket
+func (client *Client) Write() {
+	defer client.Socket.Close()
+	var err error
+	for msg := range client.Send {
+		err = client.Socket.WriteJSON(msg)
+		if err != nil {
+			log.Error("Error inside client write ", err)
+		}
+	}
+}
+
+type Hosts struct {
+	Config *config.Config
+	// Connections : hostname mapped to connection instances to reuse
+	// across metrics
+	mu      sync.Mutex
+	Drivers map[string]*driver.Driver
+	Client  chan *Client
+	Start   chan bool
+}
+
+func (hosts *Hosts) getDriver(address string) *driver.Driver {
+	hosts.mu.Lock()
+	defer hosts.mu.Unlock()
+	return hosts.Drivers[address]
+}
+
+func (hosts *Hosts) resetDriver(host config.Host) {
+	hosts.mu.Lock()
+	defer hosts.mu.Unlock()
+	hostDriver := host.Connection.ToDriver()
+	hosts.Drivers[host.Address] = &hostDriver
+}
+
+func (hosts *Hosts) sendMetric(host config.Host, client *Client) {
+	if hosts.getDriver(host.Address) == nil {
+		hosts.resetDriver(host)
+	}
+	for _, metric := range config.GetDashboardInfoConfig(hosts.Config).Metrics {
+		initializedMetric, err := inspector.Init(metric, hosts.getDriver(host.Address))
+		data, err := initializedMetric.Execute()
+		if err == nil {
+			var unmarsh interface{}
+			json.Unmarshal(data, &unmarsh)
+			message := &FullMessage{
+				Message: Message{
+					Host: host.Address,
+					Name: metric,
+					Data: unmarsh,
+				},
+				Error: false,
+			}
+			client.Send <- message
+		} else {
+			// check for error 127 which means command was not found
+			var errorContent string
+			if !strings.Contains(fmt.Sprintf("%s", err), "127") {
+				errorContent = fmt.Sprintf("Could not retrieve metric %s from driver %s with error %s, resetting connection...", metric, host.Address, err)
+			} else {
+				errorContent = fmt.Sprintf("Command %s not found on driver %s", metric, host.Address)
+			}
+			log.Error(errorContent)
+			hosts.resetDriver(host)
+			message := &FullMessage{
+				Message: errorContent,
+				Error:   true,
+			}
+			client.Send <- message
+		}
+	}
+}
+
+func (hosts *Hosts) Run() {
+	dashboardInfo := config.GetDashboardInfoConfig(hosts.Config)
+	log.Debug("In Running")
+	for {
+		select {
+		case client := <-hosts.Client:
+			for {
+				for _, host := range dashboardInfo.Hosts {
+					go hosts.sendMetric(host, client)
+				}
+				log.Infof("Delaying for %d seconds", dashboardInfo.PollInterval)
+				time.Sleep(time.Duration(dashboardInfo.PollInterval) * time.Second)
+			}
+		}
+	}
+
+}
+
+func (hosts *Hosts) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	socket, err := upgrader.Upgrade(w, req, nil)
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+	client := &Client{
+		Socket: socket,
+		Send:   make(chan *FullMessage, messageBufferSize),
+	}
+	hosts.Client <- client
+	client.Write()
+}
+
+func newHosts(cfg *config.Config) *Hosts {
+	hosts := &Hosts{
+		Config:  cfg,
+		Drivers: make(map[string]*driver.Driver),
+		Client:  make(chan *Client),
+	}
+	return hosts
+}
+
+func setHostHandler(w http.ResponseWriter, r *http.Request) {
+	b, _ := json.Marshal("Hello World")
+	w.Write(b)
+}
+
+var apiCmd = &cobra.Command{
+	Use:   "api",
+	Short: "host saido as an API on a PORT env variable, fallback to set argument",
+	Long:  ``,
+	Run: func(cmd *cobra.Command, args []string) {
+		//    server.HandleFunc("/set-hosts", SetHostHandler)
+		// FIXME: set up cfg using set-hosts endpoint
+		hosts := newHosts(cfg)
+		server.HandleFunc("/set-hosts", setHostHandler)
+		server.Handle("/metrics", hosts)
+		log.Info("listening on :", port)
+		_, err := strconv.Atoi(port)
+		if err != nil {
+			log.Fatal(err)
+		}
+		go hosts.Run()
+		loggedRouters := handlers.LoggingHandler(os.Stdout, server)
+		if err := http.ListenAndServe(":"+port, loggedRouters); err != nil {
+			log.Fatal(err)
+		}
+	},
+}
+
+func init() {
+	apiCmd.Flags().StringVarP(&port, "port", "p", "3000", "Port to run application server on")
+	rootCmd.AddCommand(apiCmd)
+}

config.example.yaml

@@ -29,4 +29,5 @@ hosts:
 
 metrics:
 - memory
-- cpu
+- tcp
+poll-interval: 30

config/config.go

@@ -4,16 +4,18 @@ import (
 	"fmt"
 	"io/ioutil"
 
+	"github.com/bisohns/saido/driver"
 	"github.com/mitchellh/mapstructure"
 	log "github.com/sirupsen/logrus"
 
 	"gopkg.in/yaml.v2"
 )
 
 type DashboardInfo struct {
-	Hosts   []Host
-	Metrics []string
-	Title   string
+	Hosts        []Host
+	Metrics      []string
+	Title        string
+	PollInterval int
 }
 
 type Connection struct {
@@ -22,6 +24,22 @@ type Connection struct {
 	Password       string `mapstructure:"password"`
 	PrivateKeyPath string `mapstructure:"private_key_path"`
 	Port           int32  `mapstructure:"port"`
+	Host           string
+}
+
+func (conn *Connection) ToDriver() driver.Driver {
+	switch conn.Type {
+	case "ssh":
+		return &driver.SSH{
+			User:            conn.Username,
+			Host:            conn.Host,
+			Port:            int(conn.Port),
+			KeyFile:         conn.PrivateKeyPath,
+			CheckKnownHosts: false,
+		}
+	default:
+		return &driver.Local{}
+	}
 }
 
 type Host struct {
@@ -31,9 +49,10 @@ type Host struct {
 }
 
 type Config struct {
-	Hosts   map[interface{}]interface{} `yaml:"hosts"`
-	Metrics []string                    `yaml:"metrics"`
-	Title   string                      `yaml:"title"`
+	Hosts        map[interface{}]interface{} `yaml:"hosts"`
+	Metrics      []string                    `yaml:"metrics"`
+	Title        string                      `yaml:"title"`
+	PollInterval int                         `yaml:"poll-interval"`
 }
 
 func LoadConfig(configPath string) *Config {
@@ -62,6 +81,7 @@ func GetDashboardInfoConfig(config *Config) *DashboardInfo {
 	for _, host := range dashboardInfo.Hosts {
 		log.Debugf("%s: %v", host.Address, host.Connection)
 	}
+	dashboardInfo.PollInterval = config.PollInterval
 	return dashboardInfo
 }
 
@@ -107,6 +127,7 @@ func parseConfig(name string, host string, group map[interface{}]interface{}, cu
 	}
 
 	if !isParent {
+		currentConn.Host = host
 		newHost := Host{
 			Address:    host,
 			Connection: currentConn,

go.mod

@@ -3,6 +3,8 @@ module github.com/bisohns/saido
 go 1.14
 
 require (
+	github.com/gorilla/handlers v1.5.1 // indirect
+	github.com/gorilla/websocket v1.5.0 // indirect
 	github.com/kr/pretty v0.2.0 // indirect
 	github.com/melbahja/goph v1.2.1
 	github.com/mitchellh/mapstructure v1.4.3

go.sum

@@ -67,6 +67,8 @@ github.com/envoyproxy/go-control-plane v0.9.9-0.20201210154907-fd9021fe5dad/go.m
 github.com/envoyproxy/go-control-plane v0.9.9-0.20210217033140-668b12f5399d/go.mod h1:cXg6YxExXjJnVBQHBLXeUAgxn2UodCpnH306RInaBQk=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
+github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
+github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
 github.com/gdamore/encoding v1.0.0 h1:+7OoQ1Bc6eTm5niUzBa0Ctsh6JbMW6Ra+YNuAtDBdko=
 github.com/gdamore/encoding v1.0.0/go.mod h1:alR0ol34c49FCSBLjhosxzcPHQbf2trDkoo5dl+VrEg=
@@ -140,6 +142,10 @@ github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
+github.com/gorilla/handlers v1.5.1 h1:9lRY6j8DEeeBT10CvO9hGW0gmky0BprnvDI5vfhUHH4=
+github.com/gorilla/handlers v1.5.1/go.mod h1:t8XrUpc4KVXb7HGyJ4/cEnwQiaxrX/hz1Zv/4g96P1Q=
+github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
+github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
 github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBtguAZLlVdkD9Q=
 github.com/hashicorp/consul/sdk v0.1.1/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8=

inspector/custom.go

@@ -1,6 +1,7 @@
 package inspector
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 
@@ -22,7 +23,7 @@ type Custom struct {
 
 // Parse : run custom parsing on output of the command
 func (i *Custom) Parse(output string) {
-	log.Debug("Parsing ouput string in Custom inspector")
+	log.Debug("Parsing output string in Custom inspector")
 	i.Values = i.createMetric(output)
 }
 
@@ -44,11 +45,13 @@ func (i Custom) driverExec() driver.Command {
 	return (*i.Driver).RunCommand
 }
 
-func (i *Custom) Execute() {
+func (i *Custom) Execute() ([]byte, error) {
 	output, err := i.driverExec()(i.Command)
 	if err == nil {
 		i.Parse(output)
+		return json.Marshal(i.Values)
 	}
+	return []byte(""), err
 }
 
 // NewCustom : Initialize a new Custom instance