-
-
Notifications
You must be signed in to change notification settings - Fork 5
/
nsrl.go
144 lines (121 loc) · 2.89 KB
/
nsrl.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
package filetrove
import (
"bufio"
"bytes"
"errors"
"github.com/schollz/progressbar/v3"
"go.etcd.io/bbolt"
"io"
"net/http"
"os"
"strings"
)
func CreateNSRLBoltDB(nsrlsourcefile string, nsrldbfile string) error {
db, err := bbolt.Open(nsrldbfile, 0600, nil)
if err != nil {
return err
}
defer db.Close()
file, err := os.Open(nsrlsourcefile)
if err != nil {
return err
}
defer file.Close()
batchSize := 100000
values := make([]string, 0, batchSize)
scanner := bufio.NewScanner(file)
for scanner.Scan() {
hash := scanner.Text()
values = append(values, hash)
if len(values) == batchSize {
err := db.Update(func(tx *bbolt.Tx) error {
bucket, err := tx.CreateBucketIfNotExists([]byte("sha1"))
if err != nil {
return err
}
// Reduce file size
bucket.FillPercent = 0.9
for _, value := range values {
err := bucket.Put([]byte(strings.ToLower(value)), []byte("true"))
if err != nil {
return err
}
}
return nil
})
if err != nil {
return err
}
values = values[:0]
}
}
if len(values) > 0 {
err := db.Update(func(tx *bbolt.Tx) error {
bucket, err := tx.CreateBucketIfNotExists([]byte("sha1"))
if err != nil {
return err
}
for _, value := range values {
err := bucket.Put([]byte(strings.ToLower(value)), []byte("true"))
if err != nil {
return err
}
}
return nil
})
if err != nil {
return err
}
}
return nil
}
// GetNSRL downloads a prepared BoltDB database file from an online storage
func GetNSRL() error {
req, err := http.NewRequest("GET", "https://download.fritz.wtf/nsrl.db", nil)
if err != nil {
return err
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
f, err := os.OpenFile("db/nsrl.db", os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
return err
}
defer f.Close()
bar := progressbar.DefaultBytes(
resp.ContentLength,
"downloading",
)
io.Copy(io.MultiWriter(f, bar), resp.Body)
return nil
}
// ChecksumNSRL checks a NSRL BoltDB's checksum that is provided with a sidecar file
func ChecksumNSRL(nsrldbfile string) {
Hashit(nsrldbfile, "blake2b-512")
}
// ConnectNSRL connects to local bbolt NSRL file
func ConnectNSRL(nsrldbfile string) (*bbolt.DB, error) {
db, err := bbolt.Open(nsrldbfile, 0600, nil)
if err != nil {
return nil, err
}
return db, nil
}
// GetValueNSRL reads bbolt database and checks if a given sha1 hash is present in the database
func GetValueNSRL(db *bbolt.DB, sha1hash []byte) (bool, error) {
var fileIsInNSRL bool
err := db.View(func(tx *bbolt.Tx) error {
b := tx.Bucket([]byte("sha1"))
if b == nil {
return errors.New("Could not connect to bucket.")
}
// the byte array translates to UTF-8 "true"
fileIsInNSRL = bytes.Equal(b.Get(sha1hash), []byte{116, 114, 117, 101})
// return nil to complete the transaction
return nil
})
return fileIsInNSRL, err
}