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
|
package db
import (
"fmt"
"gopkg.in/cheggaaa/pb.v2"
"github.com/inconshreveable/log15"
"github.com/jinzhu/gorm"
"github.com/k0kubun/pp"
"github.com/kotakanbe/go-cpe-dictionary/models"
// Required MySQL. See http://jinzhu.me/gorm/database.html#connecting-to-a-database
_ "github.com/jinzhu/gorm/dialects/mysql"
_ "github.com/jinzhu/gorm/dialects/postgres"
// Required SQLite3.
_ "github.com/jinzhu/gorm/dialects/sqlite"
)
// Supported DB dialects.
const (
dialectSqlite3 = "sqlite3"
dialectMysql = "mysql"
dialectPostgreSQL = "postgres"
)
// RDBDriver is Driver for RDB
type RDBDriver struct {
name string
conn *gorm.DB
}
// Name return db name
func (r *RDBDriver) Name() string {
return r.name
}
// NewRDB return RDB driver
func NewRDB(dbType, dbpath string, debugSQL bool) (driver *RDBDriver, err error) {
driver = &RDBDriver{
name: dbType,
}
log15.Debug("Opening DB", "db", driver.Name())
if err = driver.OpenDB(dbType, dbpath, debugSQL); err != nil {
return
}
log15.Debug("Migrating DB.", "db", driver.Name())
if err = driver.MigrateDB(); err != nil {
return
}
return
}
// OpenDB opens Database
func (r *RDBDriver) OpenDB(dbType, dbPath string, debugSQL bool) (err error) {
r.conn, err = gorm.Open(dbType, dbPath)
if err != nil {
err = fmt.Errorf("Failed to open DB. dbtype: %s, dbpath: %s, err: %s", dbType, dbPath, err)
return
}
r.conn.LogMode(debugSQL)
if r.name == dialectSqlite3 {
r.conn.Exec("PRAGMA journal_mode=WAL;")
}
return
}
// MigrateDB migrates Database
func (r *RDBDriver) MigrateDB() error {
if err := r.conn.AutoMigrate(
&models.CategorizedCpe{},
).Error; err != nil {
return fmt.Errorf("Failed to migrate. err: %s", err)
}
errMsg := "Failed to create index. err: %s"
if err := r.conn.Model(&models.CategorizedCpe{}).
AddUniqueIndex("idx_cpes_uri", "cpe_uri").Error; err != nil {
return fmt.Errorf(errMsg, err)
}
return nil
}
// InsertCpes inserts Cpe Information into DB
func (r *RDBDriver) InsertCpes(cpes []*models.CategorizedCpe) error {
insertedCpes := []string{}
bar := pb.StartNew(len(cpes))
for chunked := range chunkSlice(cpes, 100) {
tx := r.conn.Begin()
for _, c := range chunked {
bar.Increment()
// select old record.
old := models.CategorizedCpe{}
r := tx.Where(&models.CategorizedCpe{CpeURI: c.CpeURI}).First(&old)
if r.RecordNotFound() || old.ID == 0 {
if err := tx.Create(&c).Error; err != nil {
tx.Rollback()
return fmt.Errorf("Failed to insert. cve: %s, err: %s",
pp.Sprintf("%v", c),
err,
)
}
insertedCpes = append(insertedCpes, c.CpeURI)
}
}
tx.Commit()
}
bar.Finish()
log15.Info(fmt.Sprintf("Inserted %d CPEs", len(insertedCpes)))
// log.Debugf("%v", refreshedNvds)
return nil
}
// GetVendorProducts : GetVendorProducts
func (r *RDBDriver) GetVendorProducts() (vendorProducts []string, err error) {
// TODO
return nil, nil
}
// GetCpesByVendorProduct : GetCpesByVendorProduct
func (r *RDBDriver) GetCpesByVendorProduct(vendor, product string) (cpeURIs []string, err error) {
// TODO
return nil, nil
}
// CloseDB close Database
func (r *RDBDriver) CloseDB() (err error) {
if err = r.conn.Close(); err != nil {
log15.Error("Failed to close DB.", "Type", r.name, "err", err)
return
}
return
}
|
