// database storage package dbstore import ( "context" db_sql "database/sql" "embed" "fmt" _ "github.com/mattn/go-sqlite3" "github.com/pablotron/cvez/cpedict" "github.com/pablotron/cvez/cpematch" ) //go:embed sql var sqlFs embed.FS // sqlite3 backing store type DbStore struct { db *db_sql.DB } // Open database. // // This function is called by Open(). It is a separate package-private // function to make Open() easier to test. func openFull(dbType, path string) (DbStore, error) { var r DbStore // init db if db, err := db_sql.Open(dbType, path); err != nil { return r, err } else { // save handle r.db = db return r, nil } } // Open database func Open(path string) (DbStore, error) { return openFull("sqlite3", path) } // 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 } var addCpeDictionaryQueryIds = []string { "cpe/insert", "cpe/insert-title", "cpe/insert-ref", } // 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 queries, err := getQueries(addCpeDictionaryQueryIds) 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["cpe/insert"].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["cpe/insert-title"].ExecContext(ctx, id, title.Lang, title.Text) if err != nil { return err } } // add refs for _, ref := range(item.References) { _, err := sts["cpe/insert-ref"].ExecContext(ctx, id, ref.Href, ref.Text) if err != nil { return err } } } // commit changes, return result return tx.Commit() } // 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() { if sr, err := unmarshalCpeSearchRow(rows); err != nil { // return error return r, err } else { // append to results r = append(r, sr) } } // close rows // FIXME: is this correct? i am following the example from the // database/sql documentation, but it is messy and it seems // counterintuitive to close the row set and then do an additional // test for iteration errors... 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 } // query IDs used by AddCpeMatches() var addCpeMatchesQueryIds = []string { "cpe-match-insert", "cpe-match-insert-vulnerability", "cpe-match-insert-version-min", "cpe-match-insert-version-max", "cpe-match-insert-cpe", } // import CPE matches func (me DbStore) AddCpeMatches(ctx context.Context, matches cpematch.Matches) error { // lazy-init db if err := me.Init(ctx); err != nil { return err } // build query map queries, err := getQueries(addCpeMatchesQueryIds) 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 matches for _, m := range(matches.Matches) { // add cpe rs, err := sts["cpe-match/insert"].ExecContext(ctx, m.Cpe23Uri, m.Cpe22Uri) if err != nil { return err } // get last row ID id, err := rs.LastInsertId() if err != nil { return err } // add vulnerable if m.Vulnerable != nil { _, err := sts["cpe-match/insert-vulnerable"].ExecContext(ctx, id, *m.Vulnerable) if err != nil { return err } } // add version minimum if m.VersionStartIncluding != "" { _, err := sts["cpe-match/insert-versiom-min"].ExecContext(ctx, id, true, m.VersionStartIncluding) if err != nil { return err } } else if m.VersionStartExcluding != "" { _, err := sts["cpe-match/insert-versiom-min"].ExecContext(ctx, id, false, m.VersionStartExcluding) if err != nil { return err } } // add version maximum if m.VersionEndIncluding != "" { _, err := sts["cpe-match/insert-versiom-max"].ExecContext(ctx, id, true, m.VersionEndIncluding) if err != nil { return err } } else if m.VersionEndExcluding != "" { _, err := sts["cpe-match/insert-versiom-max"].ExecContext(ctx, id, false, m.VersionEndExcluding) if err != nil { return err } } // add names for _, name := range(m.Names) { _, err := sts["cpe-match/insert-name"].ExecContext(ctx, id, name.Cpe23Uri, name.Cpe22Uri) if err != nil { return err } } } // commit changes, return result return tx.Commit() }