path: root/dbstore/dbstore.go

// 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-name",
}

// 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()
}

// search CPE matches
func (me DbStore) CpeMatchSearch(
  ctx context.Context,
  match string,
) ([]string, error) {
  var r []string

  // lazy-init db
  if err := me.Init(ctx); err != nil {
    return r, err
  }

  // get query
  // FIXME: cache this?
  sql, err := getQuery("cpe-match/search.sql")
  if err != nil {
    return r, err
  }

  // exec search query
  rows, err := me.db.QueryContext(ctx, sql, match)
  if err != nil {
    return r, err
  }

  // walk results
  for rows.Next() {
    var s string
    if err := rows.Scan(&s); err != nil {
      // return error
      return r, err
    } else {
      // append to results
      r = append(r, s)
    }
  }

  // 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
}