-
Notifications
You must be signed in to change notification settings - Fork 373
/
db.go
149 lines (129 loc) · 3.13 KB
/
db.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
package db
import (
"fmt"
"math/rand"
"os"
"path"
"path/filepath"
"time"
"github.com/zan8in/afrog/v3/pkg/poc"
"github.com/zan8in/gologger"
snowflake "github.com/zan8in/pins/snowflake"
"gopkg.in/yaml.v2"
)
type Result struct {
TaskID string `json:"taskid"`
VulID string `json:"vulid"`
VulName string `json:"vulname"`
Target string `json:"target"`
FullTarget string `json:"fulltarget,omitempty"`
Severity string `json:"severity"`
Poc *poc.Poc `json:"poc,omitempty"`
Result []*PocResult `json:"pocresult,omitempty"`
Created time.Time `json:"created"`
FingerPrint any `json:"fingerprint"`
Extractor yaml.MapSlice `json:"extractor"`
}
type PocResult struct {
FullTarget string `json:"fulltarget,omitempty"`
Request string `json:"request,omitempty"`
Response string `json:"response,omitempty"`
Other Other `json:"other,omitempty"`
}
type Other struct {
Latency int64 `json:"latency,omitempty"`
}
type ResultData struct {
ID int64
TaskID string
VulID string
VulName string
Target string
FullTarget string
Severity string
Poc string
Result string
Created string
FingerPrint string
Extractor string
ResultList []PocResult
PocInfo poc.Poc
}
var (
LIMIT = "100"
DBName = "afrog"
TableName = "result"
SqliteCreate = `CREATE TABLE IF NOT EXISTS "result" (
"id" INTEGER NOT NULL DEFAULT '',
"taskid" text NOT NULL DEFAULT '',
"vulid" text NOT NULL DEFAULT '',
"vulname" text NOT NULL DEFAULT '',
"target" TEXT NOT NULL DEFAULT '',
"fulltarget" TEXT NOT NULL DEFAULT '',
"severity" TEXT NOT NULL DEFAULT '',
"poc" TEXT NOT NULL DEFAULT '',
"result" TEXT NOT NULL DEFAULT '',
"created" TEXT NOT NULL DEFAULT '',
"fingerprint" TEXT NOT NULL DEFAULT '',
"extractor" TEXT NOT NULL DEFAULT '',
PRIMARY KEY ("id")
);
CREATE INDEX "idx_search"
ON "result" (
"taskid",
"vulid",
"vulname",
"severity"
);
CREATE INDEX "idx_severity"
ON "result" (
"severity" ASC
);
CREATE INDEX "idx_taskid"
ON "result" (
"taskid" ASC
);
CREATE INDEX "idx_vulname"
ON "result" (
"vulname"
);
CREATE INDEX "idx_vulid"
ON "result" (
"vulid"
);`
TaskID string
)
var SnowFlake *snowflake.Snowflake
func init() {
TaskID = createTaskID()
if err := NewSnowFlake(); err != nil {
gologger.Fatal().Msgf("New SnowFlake failed: %v", err)
}
}
func createTaskID() string {
timestamp := time.Now().UnixNano()
source := rand.NewSource(time.Now().UnixNano())
randomGenerator := rand.New(source)
randomNum := randomGenerator.Intn(10000)
taskID := fmt.Sprintf("%d%d", timestamp, randomNum)
return taskID
}
func DbName() string {
homeDir, err := os.UserHomeDir()
if err != nil {
return ""
}
path := path.Join(homeDir, ".config", "afrog")
if err := os.MkdirAll(path, os.ModePerm); err != nil {
return ""
}
return filepath.Join(path, DBName+".db")
}
func NewSnowFlake() error {
if node, err := snowflake.NewSnowflake(1); err != nil {
return err
} else {
SnowFlake = node
return nil
}
}