internal/dbstore: working dbstore cpe search

3 files changed, 441 insertions, 1 deletions

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
+      }
+    })
+  }
+}