-
Notifications
You must be signed in to change notification settings - Fork 5
/
agent_handler.go
95 lines (76 loc) · 2.41 KB
/
agent_handler.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
package controllers
import (
"net/http"
log "github.com/sirupsen/logrus"
"github.com/google/uuid"
"github.com/gorilla/websocket"
"github.com/labstack/echo/v4"
"github.com/lachlan2k/phatcrack/api/internal/db"
"github.com/lachlan2k/phatcrack/api/internal/filerepo"
"github.com/lachlan2k/phatcrack/api/internal/fleet"
"github.com/lachlan2k/phatcrack/api/internal/util"
)
func HookAgentHandlerEndpoints(api *echo.Group) {
// NOTE: this is just for agent handling
// These endpoints are exempt from useful authz/n
api.GET("/ping", func(c echo.Context) error {
return c.String(http.StatusOK, "pong agent")
})
api.GET("/ws", handleAgentWs)
api.GET("/download-file/:id", handleAgentDownloadFile)
}
func handleAgentDownloadFile(c echo.Context) error {
fileId := c.Param("id")
if !util.AreValidUUIDs(fileId) {
return echo.ErrBadRequest
}
authKey := c.Request().Header.Get("Authorization")
if len(authKey) == 0 {
return echo.ErrUnauthorized
}
agentId, err := db.FindAgentIDByAuthKey(authKey)
if err != nil || agentId == "" {
return echo.ErrUnauthorized
}
filename, err := filerepo.GetPathToFile(uuid.MustParse(fileId))
if err != nil {
return util.ServerError("Failed to get file", err)
}
err = c.File(filename)
if err == echo.ErrNotFound {
log.WithField("file_id", fileId).WithField("agent_id", agentId).WithError(err).Warn("Agent tried to download a file that doesn't exist")
return err
}
if err != nil {
log.WithField("file_id", fileId).WithField("agent_id", agentId).WithError(err).Warn("Agent tried to download a file but encountered an error")
}
return err
}
func handleAgentWs(c echo.Context) error {
authKey := c.Request().Header.Get("Authorization")
if len(authKey) == 0 {
return echo.ErrUnauthorized
}
agentData, err := db.FindAgentByAuthKey(authKey)
if err != nil {
return echo.ErrUnauthorized
}
ws, err := (&websocket.Upgrader{}).Upgrade(c.Response(), c.Request(), nil)
if err != nil {
return util.ServerError("Couldn't upgrade websocket", err)
}
defer ws.Close()
agent := fleet.RegisterAgentFromWebsocket(ws, agentData.ID.String())
AuditLog(c, log.Fields{
"agent_id": agentData.ID.String(),
"agent_name": agentData.Name,
}, "Agent has connected and is being handled")
err = agent.Handle()
if err != nil {
log.WithError(err).WithFields(log.Fields{
"agent_id": agentData.ID.String(),
"agent_name": agentData.Name,
}).Warn("Error from agent")
}
return nil
}