diff options
Diffstat (limited to 'internal')
| -rw-r--r-- | internal/dbstore/cpesearchtype_string.go | 25 | ||||
| -rw-r--r-- | internal/dbstore/dbstore.go | 252 | ||||
| -rw-r--r-- | internal/dbstore/dbstore_test.go | 165 | 
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 +      } +    }) +  } +}  | 
