forked from gophish/gophish
/
group.go
313 lines (290 loc) · 8.26 KB
/
group.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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
package models
import (
"errors"
"fmt"
"net/mail"
"time"
"github.com/jinzhu/gorm"
)
// Group contains the fields needed for a user -> group mapping
// Groups contain 1..* Targets
type Group struct {
Id int64 `json:"id"`
UserId int64 `json:"-"`
Name string `json:"name"`
ModifiedDate time.Time `json:"modified_date"`
Targets []Target `json:"targets" sql:"-"`
}
// GroupSummaries is a struct representing the overview of Groups.
type GroupSummaries struct {
Total int64 `json:"total"`
Groups []GroupSummary `json:"groups"`
}
// GroupSummary represents a summary of the Group model. The only
// difference is that, instead of listing the Targets (which could be expensive
// for large groups), it lists the target count.
type GroupSummary struct {
Id int64 `json:"id"`
Name string `json:"name"`
ModifiedDate time.Time `json:"modified_date"`
NumTargets int64 `json:"num_targets"`
}
// GroupTarget is used for a many-to-many relationship between 1..* Groups and 1..* Targets
type GroupTarget struct {
GroupId int64 `json:"-"`
TargetId int64 `json:"-"`
}
// Target contains the fields needed for individual targets specified by the user
// Groups contain 1..* Targets, but 1 Target may belong to 1..* Groups
type Target struct {
Id int64 `json:"-"`
FirstName string `json:"first_name"`
LastName string `json:"last_name"`
Email string `json:"email"`
Position string `json:"position"`
}
// Returns the email address to use in the "To" header of the email
func (t *Target) FormatAddress() string {
addr := t.Email
if t.FirstName != "" && t.LastName != "" {
a := &mail.Address{
Name: fmt.Sprintf("%s %s", t.FirstName, t.LastName),
Address: t.Email,
}
addr = a.String()
}
return addr
}
// ErrNoEmailSpecified is thrown when no email is specified for the Target
var ErrEmailNotSpecified = errors.New("No email address specified")
// ErrGroupNameNotSpecified is thrown when a group name is not specified
var ErrGroupNameNotSpecified = errors.New("Group name not specified")
// ErrNoTargetsSpecified is thrown when no targets are specified by the user
var ErrNoTargetsSpecified = errors.New("No targets specified")
// Validate performs validation on a group given by the user
func (g *Group) Validate() error {
switch {
case g.Name == "":
return ErrGroupNameNotSpecified
case len(g.Targets) == 0:
return ErrNoTargetsSpecified
}
return nil
}
// GetGroups returns the groups owned by the given user.
func GetGroups(uid int64) ([]Group, error) {
gs := []Group{}
err := db.Where("user_id=?", uid).Find(&gs).Error
if err != nil {
Logger.Println(err)
return gs, err
}
for i, _ := range gs {
gs[i].Targets, err = GetTargets(gs[i].Id)
if err != nil {
Logger.Println(err)
}
}
return gs, nil
}
// GetGroupSummaries returns the summaries for the groups
// created by the given uid.
func GetGroupSummaries(uid int64) (GroupSummaries, error) {
gs := GroupSummaries{}
query := db.Table("groups").Where("user_id=?", uid)
err := query.Select("id, name, modified_date").Scan(&gs.Groups).Error
if err != nil {
Logger.Println(err)
return gs, err
}
for i := range gs.Groups {
query = db.Table("group_targets").Where("group_id=?", gs.Groups[i].Id)
err = query.Count(&gs.Groups[i].NumTargets).Error
if err != nil {
return gs, err
}
}
gs.Total = int64(len(gs.Groups))
return gs, nil
}
// GetGroup returns the group, if it exists, specified by the given id and user_id.
func GetGroup(id int64, uid int64) (Group, error) {
g := Group{}
err := db.Where("user_id=? and id=?", uid, id).Find(&g).Error
if err != nil {
Logger.Println(err)
return g, err
}
g.Targets, err = GetTargets(g.Id)
if err != nil {
Logger.Println(err)
}
return g, nil
}
// GetGroupSummary returns the summary for the requested group
func GetGroupSummary(id int64, uid int64) (GroupSummary, error) {
g := GroupSummary{}
query := db.Table("groups").Where("user_id=? and id=?", uid, id)
err := query.Select("id, name, modified_date").Scan(&g).Error
if err != nil {
Logger.Println(err)
return g, err
}
query = db.Table("group_targets").Where("group_id=?", id)
err = query.Count(&g.NumTargets).Error
if err != nil {
return g, err
}
return g, nil
}
// GetGroupByName returns the group, if it exists, specified by the given name and user_id.
func GetGroupByName(n string, uid int64) (Group, error) {
g := Group{}
err := db.Where("user_id=? and name=?", uid, n).Find(&g).Error
if err != nil {
Logger.Println(err)
return g, err
}
g.Targets, err = GetTargets(g.Id)
if err != nil {
Logger.Println(err)
}
return g, err
}
// PostGroup creates a new group in the database.
func PostGroup(g *Group) error {
if err := g.Validate(); err != nil {
return err
}
// Insert the group into the DB
err = db.Save(g).Error
if err != nil {
Logger.Println(err)
return err
}
for _, t := range g.Targets {
insertTargetIntoGroup(t, g.Id)
}
return nil
}
// PutGroup updates the given group if found in the database.
func PutGroup(g *Group) error {
if err := g.Validate(); err != nil {
return err
}
// Fetch group's existing targets from database.
ts := []Target{}
ts, err = GetTargets(g.Id)
if err != nil {
Logger.Printf("Error getting targets from group ID: %d", g.Id)
return err
}
// Check existing targets, removing any that are no longer in the group.
tExists := false
for _, t := range ts {
tExists = false
// Is the target still in the group?
for _, nt := range g.Targets {
if t.Email == nt.Email {
tExists = true
break
}
}
// If the target does not exist in the group any longer, we delete it
if !tExists {
err = db.Where("group_id=? and target_id=?", g.Id, t.Id).Delete(&GroupTarget{}).Error
if err != nil {
Logger.Printf("Error deleting email %s\n", t.Email)
}
}
}
// Add any targets that are not in the database yet.
for _, nt := range g.Targets {
// Check and see if the target already exists in the db
tExists = false
for _, t := range ts {
if t.Email == nt.Email {
tExists = true
nt.Id = t.Id
break
}
}
// Add target if not in database, otherwise update target information.
if !tExists {
insertTargetIntoGroup(nt, g.Id)
} else {
UpdateTarget(nt)
}
}
err = db.Save(g).Error
if err != nil {
Logger.Println(err)
return err
}
return nil
}
// DeleteGroup deletes a given group by group ID and user ID
func DeleteGroup(g *Group) error {
// Delete all the group_targets entries for this group
err := db.Where("group_id=?", g.Id).Delete(&GroupTarget{}).Error
if err != nil {
Logger.Println(err)
return err
}
// Delete the group itself
err = db.Delete(g).Error
if err != nil {
Logger.Println(err)
return err
}
return err
}
func insertTargetIntoGroup(t Target, gid int64) error {
if _, err = mail.ParseAddress(t.Email); err != nil {
Logger.Printf("Invalid email %s\n", t.Email)
return err
}
trans := db.Begin()
trans.Where(t).FirstOrCreate(&t)
if err != nil {
Logger.Printf("Error adding target: %s\n", t.Email)
return err
}
err = trans.Where("group_id=? and target_id=?", gid, t.Id).Find(&GroupTarget{}).Error
if err == gorm.ErrRecordNotFound {
err = trans.Save(&GroupTarget{GroupId: gid, TargetId: t.Id}).Error
if err != nil {
Logger.Println(err)
return err
}
}
if err != nil {
Logger.Printf("Error adding many-many mapping for %s\n", t.Email)
return err
}
err = trans.Commit().Error
if err != nil {
Logger.Printf("Error committing db changes\n")
return err
}
return nil
}
// UpdateTarget updates the given target information in the database.
func UpdateTarget(target Target) error {
targetInfo := map[string]interface{}{
"first_name": target.FirstName,
"last_name": target.LastName,
"position": target.Position,
}
err := db.Model(&target).Where("id = ?", target.Id).Updates(targetInfo).Error
if err != nil {
Logger.Printf("Error updating target information for %s\n", target.Email)
}
return err
}
// GetTargets performs a many-to-many select to get all the Targets for a Group
func GetTargets(gid int64) ([]Target, error) {
ts := []Target{}
err := db.Table("targets").Select("targets.id, targets.email, targets.first_name, targets.last_name, targets.position").Joins("left join group_targets gt ON targets.id = gt.target_id").Where("gt.group_id=?", gid).Scan(&ts).Error
return ts, err
}