-
Notifications
You must be signed in to change notification settings - Fork 28
/
ent.go
126 lines (115 loc) · 2.87 KB
/
ent.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
124
125
126
package c2test
import (
"context"
"crypto/rand"
"encoding/json"
"fmt"
"testing"
"github.com/stretchr/testify/require"
"realm.pub/tavern/internal/c2/c2pb"
"realm.pub/tavern/internal/c2/epb"
"realm.pub/tavern/internal/ent"
"realm.pub/tavern/internal/ent/beacon"
"realm.pub/tavern/internal/ent/file"
"realm.pub/tavern/internal/namegen"
)
// NewRandomBeacon creates and returns a new randomized beacon.
func NewRandomBeacon(ctx context.Context, graph *ent.Client) *ent.Beacon {
host := graph.Host.Create().
SetPlatform(c2pb.Host_PLATFORM_UNSPECIFIED).
SetIdentifier(namegen.NewComplex()).
SaveX(ctx)
return graph.Beacon.Create().
SetIdentifier(namegen.NewComplex()).
SetHost(host).
SaveX(ctx)
}
// ConvertTaskToC2PB converts an *ent.Task to it's corresponding *c2pb.Task.
func ConvertTaskToC2PB(t *testing.T, ctx context.Context, task *ent.Task) *c2pb.Task {
t.Helper()
if task == nil {
return nil
}
var params map[string]string
require.NoError(t,
json.Unmarshal(
[]byte(
task.
QueryQuest().
OnlyX(ctx).
Parameters,
),
¶ms,
),
)
var fileNames []string
files := task.
QueryQuest().
QueryTome().
QueryFiles().
Order(file.ByID()).
AllX(ctx)
for _, f := range files {
fileNames = append(fileNames, f.Name)
}
return &c2pb.Task{
Id: int64(task.ID),
Tome: &epb.Tome{
Eldritch: task.
QueryQuest().
QueryTome().
OnlyX(ctx).
Eldritch,
Parameters: params,
FileNames: fileNames,
},
QuestName: task.
QueryQuest().
OnlyX(ctx).
Name,
}
}
// NewRandomAssignedTask generates a random task and assigns it to the provided beacon identifier.
func NewRandomAssignedTask(ctx context.Context, graph *ent.Client, beaconIdentifier string) *ent.Task {
beacon := graph.Beacon.Query().
Where(
beacon.Identifier(beaconIdentifier),
).OnlyX(ctx)
bundle := graph.File.Create().
SetName(namegen.NewComplex()).
SetContent(newRandomBytes(1024)).
SaveX(ctx)
files := []*ent.File{
graph.File.Create().
SetName(namegen.NewComplex()).
SetContent(newRandomBytes(1024)).
SaveX(ctx),
graph.File.Create().
SetName(namegen.NewComplex()).
SetContent(newRandomBytes(1024)).
SaveX(ctx),
}
tome := graph.Tome.Create().
SetName(namegen.NewComplex()).
SetEldritch(fmt.Sprintf(`print("%s")`, namegen.NewComplex())).
SetDescription(string(newRandomBytes(120))).
SetAuthor("kcarretto").
SetParamDefs(`[{"name":"test-param","label":"Test","type":"string","placeholder":"Enter text..."}]`).
AddFiles(files...).
SaveX(ctx)
quest := graph.Quest.Create().
SetName(namegen.NewComplex()).
SetBundle(bundle).
SetTome(tome).
SetParameters(fmt.Sprintf(`{"test-param":"%v"}`, namegen.NewComplex())).
SaveX(ctx)
return graph.Task.Create().
SetBeacon(beacon).
SetQuest(quest).
SaveX(ctx)
}
func newRandomBytes(length int) []byte {
data := make([]byte, length)
rand.Read(data)
return data
}