-
Notifications
You must be signed in to change notification settings - Fork 43
/
domain.go
141 lines (121 loc) · 3.38 KB
/
domain.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
package yakit
import (
"context"
"github.com/jinzhu/gorm"
"github.com/yaklang/yaklang/common/log"
"github.com/yaklang/yaklang/common/schema"
"github.com/yaklang/yaklang/common/utils"
"github.com/yaklang/yaklang/common/utils/bizhelper"
"github.com/yaklang/yaklang/common/yakgrpc/ypb"
"sort"
"strings"
)
func CreateOrUpdateDomain(db *gorm.DB, hash string, i interface{}) error {
db = db.Model(&schema.Domain{})
if db := db.Where("hash = ?", hash).Assign(i).FirstOrCreate(&schema.Domain{}); db.Error != nil {
return utils.Errorf("create/update Domain failed: %s", db.Error)
}
return nil
}
func SaveDomain(db *gorm.DB, domain string, ip string) error {
host, err := GetHostByIP(db, ip)
if err != nil {
host, err = NewHost(ip)
if err != nil {
return utils.Errorf("create ip host failed: %s", err)
}
}
domains := utils.PrettifyListFromStringSplited(host.Domains, ",")
domains = append(domains, domain)
sort.Strings(domains)
host.Domains = strings.Join(domains, ",")
_ = CreateOrUpdateHost(db, host.IP, host)
d := &schema.Domain{
Domain: domain,
IPAddr: ip,
IPInteger: host.IPInteger,
}
d.FillDomainHTTPInfo()
_ = CreateOrUpdateDomain(db, d.CalcHash(), d)
return nil
}
func GetDomain(db *gorm.DB, id int64) (*schema.Domain, error) {
var req schema.Domain
if db := db.Model(&schema.Domain{}).Where("id = ?", id).First(&req); db.Error != nil {
return nil, utils.Errorf("get Domain failed: %s", db.Error)
}
return &req, nil
}
func DeleteDomainByID(db *gorm.DB, ids ...int64) error {
if len(ids) == 1 {
id := ids[0]
if db := db.Model(&schema.Domain{}).Where(
"id = ?", id,
).Unscoped().Delete(&schema.Domain{}); db.Error != nil {
return db.Error
}
return nil
}
if db = bizhelper.ExactQueryInt64ArrayOr(db, "id", ids).Unscoped().Delete(&schema.Domain{}); db.Error != nil {
return utils.Errorf("delete id(s) failed: %v", db.Error)
}
return nil
}
func FilterDomain(db *gorm.DB, params *ypb.QueryDomainsRequest) *gorm.DB {
db = bizhelper.FuzzQueryLike(db, "domain", params.GetDomainKeyword())
db = bizhelper.QueryBySpecificAddress(db, "ip_integer", params.GetNetwork())
db = bizhelper.FuzzQueryLike(db, "http_title", params.GetTitle())
return db
}
func QueryDomain(db *gorm.DB, params *ypb.QueryDomainsRequest) (*bizhelper.Paginator, []*schema.Domain, error) {
if params == nil {
return nil, nil, utils.Errorf("empty params")
}
db = db.Model(&schema.Domain{}) // .Debug()
if params.Pagination == nil {
params.Pagination = &ypb.Paging{
Page: 1,
Limit: 30,
OrderBy: "updated_at",
Order: "desc",
}
}
p := params.Pagination
db = FilterDomain(db, params)
var ret []*schema.Domain
paging, db := bizhelper.Paging(db, int(p.Page), int(p.Limit), &ret)
if db.Error != nil {
return nil, nil, utils.Errorf("paging failed: %s", db.Error)
}
return paging, ret, nil
}
func YieldDomains(db *gorm.DB, ctx context.Context) chan *schema.Domain {
outC := make(chan *schema.Domain)
go func() {
defer close(outC)
var page = 1
for {
var items []*schema.Domain
if _, b := bizhelper.NewPagination(&bizhelper.Param{
DB: db,
Page: page,
Limit: 1000,
}, &items); b.Error != nil {
log.Errorf("paging failed: %s", b.Error)
return
}
page++
for _, d := range items {
select {
case <-ctx.Done():
return
case outC <- d:
}
}
if len(items) < 1000 {
return
}
}
}()
return outC
}