/
stats.go
123 lines (109 loc) · 2.39 KB
/
stats.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package stats
import (
"encoding/json"
"errors"
"fmt"
"github.com/gmemstr/platypus/common"
"github.com/gorilla/websocket"
"io/ioutil"
"net/http"
)
type UsageStats struct {
Hostname string `json:"hostname"`
Cpu float64 `json:"cpu"`
Memory float64 `json:"memory"`
Disk float64 `json:"disk"`
Secret string `json:"secret"`
}
type Server struct {
Stats UsageStats `json:"stats"`
Online bool `json:"online"`
}
var Servers map[string] Server
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// Handles individual connections, parsing JSON data from client and writing a message.
func Handler() common.Handler {
return func(rc *common.RouterContext, w http.ResponseWriter, r *http.Request) *common.HTTPError {
hostname := ""
c, err := upgrader.Upgrade(w, r, nil)
if err != nil {
panic(err)
}
defer c.Close()
// @TODO Set up escalation process when socket closes.
c.SetCloseHandler(CloseHandler)
for {
mt, message, err := c.ReadMessage()
if err != nil {
_ = SetOffline(hostname)
break
}
stats := UsageStats{}
err = json.Unmarshal(message, &stats)
if err != nil {
break
}
hostname = stats.Hostname
secretKey, err := ioutil.ReadFile(".secret")
key := string(secretKey)
if stats.Secret != key {
_ = c.WriteMessage(mt, []byte("invalid secret key"))
_ = c.Close()
}
// Blank out secret key after comparing.
stats.Secret = ""
err = WriteStats(hostname, stats)
if err != nil {
break
}
err = c.WriteMessage(mt, []byte(""))
if err != nil {
break
}
}
return nil
}
}
func SetOffline(hostname string) error {
server, ok := Servers[hostname]
if ok {
server.Online = false
Servers[hostname] = server
}
jsonServers, err := json.MarshalIndent(Servers, "", " ")
if err != nil {
return err
}
err = ioutil.WriteFile("stats.json", jsonServers, 0644)
if err != nil {
return err
}
return nil
}
func WriteStats(hostname string, stats UsageStats) error {
server, ok := Servers[hostname]
if ok {
server.Stats = stats
server.Online = true
}
if !ok {
server = Server{
Stats: stats,
Online: true,
}
}
Servers[hostname] = server
return nil
}
// Close handler, begin chain of escalation.
func CloseHandler(code int, text string) error {
fmt.Println("ws closed")
if code != 1000 {
return errors.New("websocket closed badly")
}
return nil
}