-
Notifications
You must be signed in to change notification settings - Fork 15
/
event_new_evidence.go
143 lines (121 loc) · 5.57 KB
/
event_new_evidence.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
// Copyright 2022, Yahoo Inc.
// Licensed under the terms of the MIT. See LICENSE file in project root for terms.
package enhancementservices
import (
"context"
"fmt"
"github.com/ashirt-ops/ashirt-server/backend/database"
"github.com/ashirt-ops/ashirt-server/backend/dtos"
"github.com/ashirt-ops/ashirt-server/backend/helpers"
"github.com/ashirt-ops/ashirt-server/backend/models"
sq "github.com/Masterminds/squirrel"
)
type NewEvidencePayload struct {
Type string `json:"type" db:"type"`
EvidenceUUID string `json:"evidenceUuid" db:"uuid"`
OperationSlug string `json:"operationSlug" db:"operation_slug"`
ContentType string `json:"contentType" db:"content_type"`
GlobalVariables []dtos.GlobalVar `json:"globalVariables"`
OperationVars []dtos.OperationVar `json:"operationVariables"`
}
type ExpandedNewEvidencePayload struct {
NewEvidencePayload
EvidenceID int64 `db:"id"`
}
func getExpandedPayloadID(e ExpandedNewEvidencePayload) int64 {
return e.EvidenceID
}
// BatchBuildNewEvidencePayload creates a set of payloads for the given operation and evidence uuids.
// This function provides convenience over the alternatives: BatchBuildNewEvidencePayloadFromUUIDs and
// BatchBuildNewEvidencePayloadForAllEvidence. Note that if no evidenceUUIDs are provided,
// then all evidence is chosen for the indicated operation.
func BatchBuildNewEvidencePayload(ctx context.Context, db database.ConnectionProxy, operationID int64, evidenceUUIDs []string) ([]ExpandedNewEvidencePayload, error) {
if len(evidenceUUIDs) == 0 {
return BatchBuildNewEvidencePayloadForAllEvidence(ctx, db, operationID)
} else {
return BatchBuildNewEvidencePayloadFromUUIDs(ctx, db, operationID, evidenceUUIDs)
}
}
// BatchBuildNewEvidencePayloadFromUUIDs creates a set of payloads, ordered by evidence ID, for the given operationID and evidenceUUIDs.
// If the list of uuids is empty, then _no payloads will be returned_.
// Also see BatchBuildNewEvidencePayload, which allows for getting all evidence if no uuids are specified
func BatchBuildNewEvidencePayloadFromUUIDs(ctx context.Context, db database.ConnectionProxy, operationID int64, evidenceUUIDs []string) ([]ExpandedNewEvidencePayload, error) {
return batchBuildNewEvidencePayloadSpecial(ctx, db, func(tx database.ConnectionProxy) ([]models.Evidence, error) {
return database.GetEvidenceFromUUIDs(tx, operationID, evidenceUUIDs)
})
}
// BatchBuildNewEvidencePayloadForAllEvidence creates a set of payloads, ordered by evidence ID, for all evidence in an operation.
// Also see BatchBuildNewEvidencePayload, which allows for specifying a subset of evidence uuids
func BatchBuildNewEvidencePayloadForAllEvidence(ctx context.Context, db database.ConnectionProxy, operationID int64) ([]ExpandedNewEvidencePayload, error) {
return batchBuildNewEvidencePayloadSpecial(ctx, db, func(tx database.ConnectionProxy) ([]models.Evidence, error) {
return database.GetAllEvidenceForOperation(tx, operationID)
})
}
func batchBuildNewEvidencePayloadSpecial(ctx context.Context, db database.ConnectionProxy,
fetch func(tx database.ConnectionProxy) ([]models.Evidence, error),
) ([]ExpandedNewEvidencePayload, error) {
var payloads []ExpandedNewEvidencePayload
err := db.WithTx(ctx, func(tx *database.Transactable) {
evidence, _ := fetch(tx)
ids := helpers.Map(evidence, database.EvidenceToID)
payloads, _ = batchBuildNewEvidencePayloadFromIDs(tx, ids)
})
return payloads, err
}
// batchBuildNewEvidencePayloadFromIDs builds a payload by getting all of the necessary details in bulk.
// Note: this relies on the ordering of evidenceIDs. No particular order is required as input,
// but the result is ordered by evidenceID, in ASC order.
func batchBuildNewEvidencePayloadFromIDs(db database.ConnectionProxy, evidenceIDs []int64) ([]ExpandedNewEvidencePayload, error) {
var payloads []ExpandedNewEvidencePayload
var globalVariables []models.GlobalVar
err := db.Select(&globalVariables, sq.Select("name", "value").From("global_vars"))
if err != nil {
return nil, fmt.Errorf("unable to gather global variables for worker")
}
var globalVariablesDTO []dtos.GlobalVar
for _, v := range globalVariables {
globalVariablesDTO = append(globalVariablesDTO, dtos.GlobalVar{
Name: v.Name,
Value: v.Value,
})
}
err = db.Select(&payloads, sq.Select(
"e.id AS id",
"e.uuid AS uuid",
"e.content_type",
"slug AS operation_slug",
"'evidence_created' AS type", // hardcode in the type so we don't have to edit each entry manually
).
From("evidence e").
LeftJoin("operations o ON e.operation_id = o.id").
Where(sq.Eq{"e.id": evidenceIDs}).
OrderBy(`e.id`),
)
if err != nil {
return nil, fmt.Errorf("unable to gather evidence data for worker")
}
for i := range payloads {
var operationVariables []models.OperationVar
err := db.Select(&operationVariables,
sq.Select("ov.name", "ov.value", "ov.slug").
From("operation_vars ov").
Join("var_operation_map vom ON ov.id = vom.var_id").
Join("operations o ON o.id = vom.operation_id").
Where(sq.Eq{"o.slug": payloads[i].OperationSlug}))
if err != nil {
return nil, fmt.Errorf("unable to gather operation variables for worker")
}
var operationVariablesDTO []dtos.OperationVar
for _, v := range operationVariables {
operationVariablesDTO = append(operationVariablesDTO, dtos.OperationVar{
Name: v.Name,
Value: v.Value,
VarSlug: v.Slug,
OperationSlug: payloads[i].OperationSlug,
})
}
payloads[i].OperationVars = operationVariablesDTO
payloads[i].GlobalVariables = globalVariablesDTO
}
return payloads, nil
}