-
Notifications
You must be signed in to change notification settings - Fork 42
/
server.go
105 lines (88 loc) · 3.14 KB
/
server.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
package c2
import (
"context"
"fmt"
"github.com/kcarretto/paragon/ent"
"github.com/kcarretto/paragon/graphql/models"
"github.com/kcarretto/paragon/pkg/agent/transport"
"go.uber.org/zap"
)
//go:generate mockgen -destination=mockgen_teamserver_test.go -package=c2_test github.com/kcarretto/paragon/pkg/agent/c2 Teamserver
//go:generate mockgen -destination=mockgen_transport_test.go -package=c2_test github.com/kcarretto/paragon/pkg/agent/transport ServerMessageWriter
// Teamserver provides client methods used to interact with a teamserver.
type Teamserver interface {
ClaimTasks(context.Context, models.ClaimTasksRequest) ([]*ent.Task, error)
SubmitTaskResult(context.Context, models.SubmitTaskResultRequest) error
}
// Server manages communication with agents. Upon check-in, the server will claim and respond with
// any tasks available for the agent.
type Server struct {
Teamserver
Log *zap.Logger
}
// WriteAgentMessage is a transport-agnostic method for handling communications from an agent.
func (srv Server) WriteAgentMessage(ctx context.Context, w transport.ServerMessageWriter, msg transport.AgentMessage) error {
var srvMsg transport.ServerMessage
defer func() {
w.WriteServerMessage(ctx, srvMsg)
}()
// Submit task results
srv.submitResults(ctx, msg.Results...)
// Determine target criteria based on reported agent metadata
target, err := resolveTarget(msg.Metadata)
if err != nil {
return err
}
// Claim tasks for the agent
tasks, err := srv.ClaimTasks(ctx, target)
if err != nil {
return fmt.Errorf("failed to claim tasks from teamserver: %w", err)
}
srvMsg.Tasks = convertTasks(tasks)
return nil
}
// Submit task results
func (srv Server) submitResults(ctx context.Context, results ...*transport.TaskResult) {
for _, result := range results {
if result == nil {
continue
}
if err := srv.SubmitTaskResult(ctx, models.SubmitTaskResultRequest{
ID: int(result.Id),
Output: &result.Output,
Error: &result.Error,
ExecStartTime: result.CoerceStartTime(),
ExecStopTime: result.CoerceStopTime(),
}); err != nil {
srv.Log.Error("failed to submit task result", zap.Error(err), zap.Int64("task_id", result.Id))
}
}
}
// convertTasks coerces an array of ent.Task to an array of transport.Task
func convertTasks(models []*ent.Task) (tasks []*transport.Task) {
for _, task := range models {
if task == nil {
continue
}
tasks = append(tasks, &transport.Task{
Id: int64(task.ID),
Content: task.Content,
})
}
return tasks
}
// resolveTarget determines parameters necessary to claim tasks based on received agent metadata.
func resolveTarget(metadata *transport.AgentMetadata) (models.ClaimTasksRequest, error) {
// Ensure Agent provided metadata to resolve a target
if metadata == nil {
return models.ClaimTasksRequest{}, fmt.Errorf("agent provided no valid metadata for target resolution")
}
// Otherwise combine Primary IP and Primary MAC addresses
return models.ClaimTasksRequest{
MachineUUID: &metadata.MachineUUID,
SessionID: &metadata.SessionID,
Hostname: &metadata.Hostname,
PrimaryIP: &metadata.PrimaryIP,
PrimaryMac: &metadata.PrimaryMAC,
}, nil
}