diff options
Diffstat (limited to 'dbstore/dbstore.go')
| -rw-r--r-- | dbstore/dbstore.go | 255 |
1 files changed, 255 insertions, 0 deletions
diff --git a/dbstore/dbstore.go b/dbstore/dbstore.go new file mode 100644 index 0000000..4df9fb9 --- /dev/null +++ b/dbstore/dbstore.go @@ -0,0 +1,255 @@ +// database storage +package dbstore + +import ( + "context" + db_sql "database/sql" + "embed" + "encoding/json" + "fmt" + _ "github.com/mattn/go-sqlite3" + "github.com/pablotron/cvez/cpedict" +) + +//go:embed sql +var sqlFs embed.FS + +type DbStore struct { + db *db_sql.DB +} + +// open database +func Open(path string) (DbStore, error) { + var r DbStore + // init db + if db, err := db_sql.Open("sqlite3", path); err != nil { + return r, err + } else { + r.db = db + return r, nil + } +} + +// initialized database version +const initDbVersion = 314159 + +func (me DbStore) isInitialized(ctx context.Context) (bool, error) { + sql := "PRAGMA user_version;" + // get version + var version int32 + if err := me.db.QueryRowContext(ctx, sql).Scan(&version); err != nil { + return false, err + } + + // return result + return (version == initDbVersion), nil +} + +// initialize database +func (me DbStore) Init(ctx context.Context) error { + if inited, err := me.isInitialized(ctx); err != nil { + return err + } else if inited { + // already initialized, return success + return nil + } + + // read init query + sql, err := sqlFs.ReadFile("sql/init.sql") + if err != nil { + return err + } + + // exec init query, return result + _, err = me.db.ExecContext(ctx, string(sql)) + return err +} + +// get single query from embedded filesystem +func getQuery(id string) (string, error) { + // read query + if data, err := sqlFs.ReadFile(fmt.Sprintf("sql/%s.sql", id)); err != nil { + return "", err + } else { + // return query + return string(data), nil + } +} + +// return query map +func getQueries(ids []string) (map[string]string, error) { + r := make(map[string]string) + + for _, id := range(ids) { + // read query + if sql, err := getQuery(id); err != nil { + return r, fmt.Errorf("%s: %s", id, err.Error()) + } else { + // save query + r[id] = sql + } + } + + // return success + return r, nil +} + +// import CPE dictionary +func (me DbStore) AddCpeDictionary(ctx context.Context, dict cpedict.Dictionary) error { + // lazy-init db + if err := me.Init(ctx); err != nil { + return err + } + + // build query map + queryIds := []string { "insert-cpe", "insert-title", "insert-ref" } + queries, err := getQueries(queryIds) + if err != nil { + return err + } + + // begin context + tx, err := me.db.BeginTx(ctx, nil) + if err != nil { + return err + } + + // build statements + sts := make(map[string]*db_sql.Stmt) + for id, sql := range(queries) { + if st, err := tx.PrepareContext(ctx, sql); err != nil { + return err + } else { + sts[id] = st + defer sts[id].Close() + } + } + + // add items + for _, item := range(dict.Items) { + // add cpe + rs, err := sts["insert-cpe"].ExecContext(ctx, item.CpeUri, item.Cpe23Item.Name) + if err != nil { + return err + } + + // get last row ID + id, err := rs.LastInsertId() + if err != nil { + return err + } + + // add titles + for _, title := range(item.Titles) { + _, err := sts["insert-title"].ExecContext(ctx, id, title.Lang, title.Text) + if err != nil { + return err + } + } + + // add refs + for _, ref := range(item.References) { + _, err := sts["insert-ref"].ExecContext(ctx, id, ref.Href, ref.Text) + if err != nil { + return err + } + } + } + + // commit changes, return result + return tx.Commit() +} + +//go:generate stringer -linecomment -type=CpeSearchType + +// CPE search type +type CpeSearchType byte + +const ( + CpeSearchAll CpeSearchType = iota //cpe-search-all + CpeSearchTitle // cpe-search-title + CpeSearchRef // cpe-search-ref +) + +// title search result +type CpeSearchRow struct { + // Database CPE ID + CpeId int64 `json:"cpe_id"` + + // v2.3 formatting string + Cpe23 string `json:"cpe23"` + + // titles + Titles []cpedict.Title `json:"titles"` + + // references + Refs []cpedict.Reference `json:"refs"` + + // search result rank + Rank float32 `json:"rank"` +} + +// search CPEs +func (me DbStore) CpeSearch( + ctx context.Context, + searchType CpeSearchType, + s string, +) ([]CpeSearchRow, error) { + var r []CpeSearchRow + + // lazy-init db + if err := me.Init(ctx); err != nil { + return r, err + } + + // get query + sql, err := getQuery(searchType.String()) + if err != nil { + return r, err + } + + // exec search query + rows, err := me.db.QueryContext(ctx, sql, db_sql.Named("q", s)) + if err != nil { + return r, err + } + + // walk results + for rows.Next() { + var sr CpeSearchRow + var titles string + var refs string + + // get row values + err = rows.Scan(&sr.CpeId, &sr.Cpe23, &titles, &refs, &sr.Rank) + if err != nil { + return r, err + } + + // unmarshal titles + if err = json.Unmarshal([]byte(titles), &sr.Titles); err != nil { + return r, err + } + + // unmarshal refs + if err = json.Unmarshal([]byte(refs), &sr.Refs); err != nil { + return r, err + } + + // append to results + r = append(r, sr) + } + + // close rows + if err = rows.Close(); err != nil { + return r, err + } + + // check for iteration errors + if err = rows.Err(); err != nil { + return r, err + } + + // return success + return r, nil +} |