Skip to content

Commit

Permalink
[wip]Adding server for websocket connections
Browse files Browse the repository at this point in the history
  • Loading branch information
deven96 committed Oct 31, 2022
1 parent 0956734 commit cb4bf03
Show file tree
Hide file tree
Showing 16 changed files with 345 additions and 49 deletions.
190 changes: 190 additions & 0 deletions cmd/api.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,190 @@
/*
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"
"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
Drivers map[string]*driver.Driver
Client chan *Client
Start chan bool
}

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 {
for _, metric := range dashboardInfo.Metrics {
if hosts.Drivers[host.Address] == nil {
hostDriver := host.Connection.ToDriver()
hosts.Drivers[host.Address] = &hostDriver
}
initializedMetric, err := inspector.Init(metric, hosts.Drivers[host.Address])
log.Debugf("%#v, %#v, %#v", metric, client, hosts.Drivers[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 {
errorContent := fmt.Sprintf("Could not retrieve metric %s from driver %s with error %s, resetting connection...", metric, host.Address, err)
log.Error(errorContent)
hosts.Drivers[host.Address] = nil
message := &FullMessage{
Message: errorContent,
Error: true,
}
client.Send <- message
}
}
}
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)
}
3 changes: 2 additions & 1 deletion config.example.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -29,4 +29,5 @@ hosts:

metrics:
- memory
- cpu
- tcp
poll-interval: 30
33 changes: 27 additions & 6 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down
7 changes: 5 additions & 2 deletions inspector/custom.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package inspector

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

Expand All @@ -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)
}

Expand All @@ -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
Expand Down
13 changes: 9 additions & 4 deletions inspector/disk.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package inspector

import (
"encoding/json"
"errors"
"fmt"
"strconv"
Expand Down Expand Up @@ -58,7 +59,7 @@ For Darwin it looks something like
*/
func (i *DF) Parse(output string) {
var values []DFMetrics
log.Debug("Parsing ouput string in DF inspector")
log.Debug("Parsing output string in DF inspector")
lines := strings.Split(output, "\n")
for index, line := range lines {
// skip title line
Expand Down Expand Up @@ -114,11 +115,13 @@ func (i DF) driverExec() driver.Command {
return (*i.Driver).RunCommand
}

func (i *DF) Execute() {
func (i *DF) Execute() ([]byte, error) {
output, err := i.driverExec()(i.Command)
if err == nil {
i.Parse(output)
return json.Marshal(i.Values)
}
return []byte(""), err
}

// DFWin: parse `wmic logicaldisk` to satisfy Inspector interface
Expand All @@ -142,7 +145,7 @@ IMANI,C:,3,191980253184,,288303964160,OS
*/
func (i *DFWin) Parse(output string) {
var values []DFMetrics
log.Debug("Parsing ouput string in DF inspector")
log.Debug("Parsing output string in DF inspector")
lineChar := "\r"
output = strings.TrimPrefix(output, lineChar)
output = strings.TrimSuffix(output, lineChar)
Expand Down Expand Up @@ -197,11 +200,13 @@ func (i DFWin) driverExec() driver.Command {
return (*i.Driver).RunCommand
}

func (i *DFWin) Execute() {
func (i *DFWin) Execute() ([]byte, error) {
output, err := i.driverExec()(i.Command)
if err == nil {
i.Parse(output)
return json.Marshal(i.Values)
}
return []byte(""), err
}

// NewDF : Initialize a new DF instance
Expand Down
Loading

0 comments on commit cb4bf03

Please sign in to comment.