From 56cc399430d127cb628b3abe84bc653d2b0ce59b Mon Sep 17 00:00:00 2001 From: Paul Duncan Date: Sat, 5 Feb 2022 02:36:07 -0500 Subject: internal/dbstore: working dbstore cpe search --- internal/dbstore/cpesearchtype_string.go | 25 +++ internal/dbstore/dbstore.go | 252 +++++++++++++++++++++++++++++++ internal/dbstore/dbstore_test.go | 165 +++++++++++++++++++- 3 files changed, 441 insertions(+), 1 deletion(-) create mode 100644 internal/dbstore/cpesearchtype_string.go (limited to 'internal') diff --git a/internal/dbstore/cpesearchtype_string.go b/internal/dbstore/cpesearchtype_string.go new file mode 100644 index 0000000..1d4826e --- /dev/null +++ b/internal/dbstore/cpesearchtype_string.go @@ -0,0 +1,25 @@ +// Code generated by "stringer -linecomment -type=CpeSearchType"; DO NOT EDIT. + +package dbstore + +import "strconv" + +func _() { + // An "invalid array index" compiler error signifies that the constant values have changed. + // Re-run the stringer command to generate them again. + var x [1]struct{} + _ = x[CpeSearchAll-0] + _ = x[CpeSearchTitle-1] + _ = x[CpeSearchRef-2] +} + +const _CpeSearchType_name = "cpe-search-allcpe-search-titlecpe-search-ref" + +var _CpeSearchType_index = [...]uint8{0, 14, 30, 44} + +func (i CpeSearchType) String() string { + if i >= CpeSearchType(len(_CpeSearchType_index)-1) { + return "CpeSearchType(" + strconv.FormatInt(int64(i), 10) + ")" + } + return _CpeSearchType_name[_CpeSearchType_index[i]:_CpeSearchType_index[i+1]] +} diff --git a/internal/dbstore/dbstore.go b/internal/dbstore/dbstore.go index 416d672..4df9fb9 100644 --- a/internal/dbstore/dbstore.go +++ b/internal/dbstore/dbstore.go @@ -1,3 +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 +} diff --git a/internal/dbstore/dbstore_test.go b/internal/dbstore/dbstore_test.go index 974573d..f957146 100644 --- a/internal/dbstore/dbstore_test.go +++ b/internal/dbstore/dbstore_test.go @@ -2,6 +2,7 @@ package dbstore import ( "compress/gzip" + "context" db_sql "database/sql" "encoding/xml" "embed" @@ -11,6 +12,7 @@ import ( "github.com/pablotron/cvez/cpedict" io_fs "io/fs" "os" + "reflect" "testing" ) @@ -67,7 +69,7 @@ func getTestQueries() (map[string]string, error) { return r, nil } -func TestSimple(t *testing.T) { +func ignoreTestSimple(t *testing.T) { testDbPath := "./testdata/foo.db" // get queries queries, err := getTestQueries() @@ -170,3 +172,164 @@ func TestSimple(t *testing.T) { return } } + +// remove file if it exists +func removeFile(path string) error { + // remove file + err := os.Remove(path) + if err != nil && errors.Is(err, io_fs.ErrNotExist) { + return nil + } + + return err +} + +func createTestDb(ctx context.Context, path string) (DbStore, error) { + // remove existing file + if err := removeFile(path); err != nil { + return DbStore{}, err + } + + // open db + return Open(path) +} + +func seedTestDb(ctx context.Context, db DbStore) error { + // load test CPEs + dict, err := getTestDictionary("testdata/test-0.xml.gz") + if err != nil { + return err + } + + // add cpe dictionary + return db.AddCpeDictionary(ctx, dict) + + // TODO: seed with other data +} + +func TestOpen(t *testing.T) { + path := "./testdata/test-new.db" + ctx := context.Background() + + if _, err := createTestDb(ctx, path); err != nil { + t.Error(err) + return + } +} + +func TestAddCpeDictionary(t *testing.T) { + path := "./testdata/test-addcpedict.db" + ctx := context.Background() + + // create db + db, err := createTestDb(ctx, path) + if err != nil { + t.Error(err) + return + } + + // load test CPEs + dict, err := getTestDictionary("testdata/test-0.xml.gz") + if err != nil { + t.Error(err) + return + } + + // add cpe dictionary + if err := db.AddCpeDictionary(ctx, dict); err != nil { + t.Error(err) + return + } +} + +// sqlite> select a.cpe23 from cpes a join (select cpe_id, min(rank) as rank from cpe_fts_all where cpe_fts_all match 'advisory' group by cpe_id) b on (b.cpe_id = a.cpe_id) order by b.rank; +// sqlite> select a.cpe23 from cpes a join (select cpe_id, min(rank) as rank from cpe_fts_all where cpe_fts_all match 'advisory AND book' group by cpe_id) b on (b.cpe_id = a.cpe_id) order by b.rank; +// cpe:2.3:a:\$0.99_kindle_books_project:\$0.99_kindle_books:6:*:*:*:*:android:*:* +// +// sqlite> select c.cpe_id, c.cpe23, a.rank from cpe_titles_fts a join cpe_titles b on (b.cpe_title_id = a.rowid) join cpes c on (c.cpe_id = b.cpe_id) where cpe_titles_fts match 'project' order by a.rank; +// 2|cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:-:*:*:*:*:node.js:*:*|-0.775759508773217 +// 3|cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:0.1.0:*:*:*:*:node.js:*:*|-0.66983333682734 +// 4|cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:0.2.0:*:*:*:*:node.js:*:*|-0.66983333682734 +// 5|cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:0.2.1:*:*:*:*:node.js:*:*|-0.66983333682734 +// 1|cpe:2.3:a:\$0.99_kindle_books_project:\$0.99_kindle_books:6:*:*:*:*:android:*:*|-0.545655647541265 +// +// sqlite> select a.cpe23 from cpes a join (select cpe_id, min(rank) as rank from cpe_fts_refs where cpe_fts_refs match 'advisory' group by cpe_id) b on (b.cpe_id = a.cpe_id) order by b.rank; +// cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:-:*:*:*:*:node.js:*:* +// cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:0.1.0:*:*:*:*:node.js:*:* +// cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:0.2.0:*:*:*:*:node.js:*:* +// cpe:2.3:a:\@thi.ng\/egf_project:\@thi.ng\/egf:0.2.1:*:*:*:*:node.js:*:* +// cpe:2.3:a:360totalsecurity:360_total_security:12.1.0.1005:*:*:*:*:*:*:* +// cpe:2.3:a:\$0.99_kindle_books_project:\$0.99_kindle_books:6:*:*:*:*:android:*:* + + +func TestCpeSearch(t *testing.T) { + path := "./testdata/test-search.db" + ctx := context.Background() + + tests := []struct { + t CpeSearchType // search type + q string // query string + exp []string // expected search results (cpe23s) + } {{ + t: CpeSearchAll, + q: "advisory AND book", + exp: []string { + "cpe:2.3:a:\\$0.99_kindle_books_project:\\$0.99_kindle_books:6:*:*:*:*:android:*:*", + }, + }, { + t: CpeSearchTitle, + q: "project", + exp: []string { + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:-:*:*:*:*:node.js:*:*", + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:0.1.0:*:*:*:*:node.js:*:*", + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:0.2.0:*:*:*:*:node.js:*:*", + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:0.2.1:*:*:*:*:node.js:*:*", + "cpe:2.3:a:\\$0.99_kindle_books_project:\\$0.99_kindle_books:6:*:*:*:*:android:*:*", + }, + }, { + t: CpeSearchRef, + q: "advisory", + exp: []string { + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:-:*:*:*:*:node.js:*:*", + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:0.1.0:*:*:*:*:node.js:*:*", + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:0.2.0:*:*:*:*:node.js:*:*", + "cpe:2.3:a:\\@thi.ng\\/egf_project:\\@thi.ng\\/egf:0.2.1:*:*:*:*:node.js:*:*", + "cpe:2.3:a:360totalsecurity:360_total_security:12.1.0.1005:*:*:*:*:*:*:*", + "cpe:2.3:a:\\$0.99_kindle_books_project:\\$0.99_kindle_books:6:*:*:*:*:android:*:*", + }, + }} + + // create db + db, err := createTestDb(ctx, path) + if err != nil { + t.Error(err) + return + } + + // seed test database + if err = seedTestDb(ctx, db); err != nil { + t.Error(err) + return + } + + for _, test := range(tests) { + t.Run(test.t.String(), func(t *testing.T) { + rows, err := db.CpeSearch(ctx, test.t, test.q) + if err != nil { + t.Error(err) + return + } + + // build ids + got := make([]string, len(rows)) + for i, row := range(rows) { + got[i] = row.Cpe23 + } + + if !reflect.DeepEqual(got, test.exp) { + t.Errorf("got \"%v\", exp \"%v\"", got, test.exp) + return + } + }) + } +} -- cgit v1.2.3