// database storage
package dbstore

import (
  "context"
  db_sql "database/sql"
  "encoding/json"
  "fmt"
  _ "github.com/mattn/go-sqlite3"
  "github.com/pablotron/cvez/cpedict"
  "github.com/pablotron/cvez/cpematch"
  nvd_feed "github.com/pablotron/cvez/feed"
)

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

// enable foreign keys
func (me DbStore) enableForeignKeys(ctx context.Context) error {
  _, err := me.db.ExecContext(ctx, "PRAGMA foreign_keys = true;")
  return err
}

// 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, enable foreign keys
    return me.enableForeignKeys(ctx)
  }

  // read init query
  if sql, err := getQuery("init"); err != nil {
    return err
  } else {
    // exec init query, return result
    _, err = me.db.ExecContext(ctx, sql)
    return err
  }
}

// Begin new transaction and create prepared statements.
func (me DbStore) Begin(ctx context.Context, queryIds []string) (Tx, error) {
  return newTx(ctx, me.db, queryIds)
}

// Create a transaction, pass it to callback, then commit the transaction
// if the callback returns success and rollback the transaction if the
// callback returns an error.
func (me DbStore) Tx(ctx context.Context, queryIds []string, fn func(Tx) error) error {
  // create transaction
  tx, err := me.Begin(ctx, queryIds)
  if err != nil {
    return err
  }

  if err := fn(tx); err != nil {
    // rollback
    if rb_err := tx.Rollback(); rb_err != nil {
      return rb_err
    }

    // return error
    return err
  } else {
    // commit transaction
    return tx.Commit()
  }
}

// Execute query and invoke callback with each row of result.
func (me DbStore) Query(
  ctx context.Context,
  queryId string,
  args []interface{},
  fn func(*db_sql.Rows) error,
) error {
  // get query
  sql, err := getQuery(queryId)
  if err != nil {
    return err
  }

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

  // walk results
  for rows.Next() {
    if err = fn(rows); err != nil {
      return err
    }
  }

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

  // check for iteration errors
  if err = rows.Err(); err != nil {
    return err
  }

  // return success
  return nil
}

// Execute query and invoke callback with each row of result.
func (me DbStore) QueryRow(
  ctx context.Context,
  queryId string,
  args []interface{},
  fn func (*db_sql.Row) error,
) error {
  // get query
  sql, err := getQuery(queryId)
  if err != nil {
    return err
  }

  // exec query, invoke callback, return result
  return fn(me.db.QueryRowContext(ctx, sql, args...))
}

// queries used by AddCpeDictionary()
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
  }

  return me.Tx(ctx, addCpeDictionaryQueryIds, func(tx Tx) error {
    // add items
    for _, item := range(dict.Items) {
      // add cpe
      rs, err := tx.Exec(ctx, "cpe/insert", 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 := tx.Exec(ctx, "cpe/insert-title", id, title.Lang, title.Text)
        if err != nil {
          return err
        }
      }

      // add refs
      for _, ref := range(item.References) {
        _, err := tx.Exec(ctx, "cpe/insert-ref", id, ref.Href, ref.Text)
        if err != nil {
          return err
        }
      }
    }

    // return success
    return nil
  })
}

// 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/exec search query
  err := me.Query(ctx, searchType.String(), []interface{} {
    db_sql.Named("q", s),
  }, func(rows *db_sql.Rows) error {
    if sr, err := unmarshalCpeSearchRow(rows); err != nil {
      // return error
      return err
    } else {
      // append to results
      r = append(r, sr)
      return nil
    }
  })

  // return results
  return r, err
}

// query IDs used by AddCpeMatches()
var addCpeMatchesQueryIds = []string {
  "cpe-match/insert",
  "cpe-match/insert-vulnerable",
  "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
  }

  // begin transaction
  return me.Tx(ctx, addCpeMatchesQueryIds, func(tx Tx) error {
    // add matches
    for _, m := range(matches.Matches) {
      // add cpe
      rs, err := tx.Exec(ctx, "cpe-match/insert", 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 := tx.Exec(ctx, "cpe-match/insert-vulnerable", id, *m.Vulnerable)
        if  err != nil {
          return err
        }
      }

      // add version minimum
      if m.VersionStartIncluding != "" && m.VersionStartExcluding != "" {
        return fmt.Errorf("cannot specify both VersionStartIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionStartIncluding, m.VersionStartExcluding)
      } else if m.VersionStartIncluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, true, m.VersionStartIncluding)
        if  err != nil {
          return err
        }
      } else if m.VersionStartExcluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, false, m.VersionStartExcluding)
        if  err != nil {
          return err
        }
      }

      // add version maximum
      if m.VersionEndIncluding != "" && m.VersionEndExcluding != "" {
        return fmt.Errorf("cannot specify both VersionEndIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionEndIncluding, m.VersionEndExcluding)
      } else if m.VersionEndIncluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, true, m.VersionEndIncluding)
        if  err != nil {
          return err
        }
      } else if m.VersionEndExcluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, false, m.VersionEndExcluding)
        if  err != nil {
          return err
        }
      }

      // add names
      for _, name := range(m.Names) {
        _, err := tx.Exec(ctx, "cpe-match/insert-name", id, name.Cpe23Uri, name.Cpe22Uri)
        if err != nil {
          return err
        }
      }
    }

    // return success
    return nil
  })
}

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

  // exec search query
  err := me.Query(ctx, "cpe-match/search", []interface{} {
    match,
  }, func(rows *db_sql.Rows) error {
    var s string
    if err := rows.Scan(&s); err != nil {
      // return error
      return err
    } else {
      // append to results
      r = append(r, s)
      return nil
    }
  })

  // return result
  return r, err
}

// Add description.
func (me DbStore) addDescriptions(ctx context.Context, tx Tx, ds []nvd_feed.Description) ([]int64, error) {
  r := make([]int64, len(ds))

  for i, d := range(ds) {
    // add description
    rs, err := tx.Exec(ctx, "feed/insert-desc", d.Lang, d.Value)
    if err != nil {
      return r, err
    }

    // return id and error
    id, err := rs.LastInsertId()
    if err != nil {
      return r, err
    }

    // add to results
    r[i] = id
  }

  // return description IDs
  return r, nil
}

// Add assigner if they are not already present in the database
func (me DbStore) addAssigner(ctx context.Context, tx Tx, itemId int64, assigner string) error {
  _, err := tx.Exec(ctx, "feed/insert-assigner", db_sql.Named("assigner", assigner))
  return err
}

// Add CVE.
func (me DbStore) addCve(ctx context.Context, tx Tx, itemId int64, cve nvd_feed.Cve) error {
  // cve ID
  var cveId int64

  // check data type, data format, and data version
  if err := checkNvdData(cve.DataType, cve.DataFormat, cve.DataVersion); err != nil {
    return err
  }

  // add cve assigner
  if err := me.addAssigner(ctx, tx, itemId, cve.Metadata.Assigner); err != nil {
    return err
  }

  // add cve
  nvdId := cve.Metadata.Id
  rs, err := tx.Exec(ctx, "feed/insert-cve", itemId, nvdId.Year(), nvdId.Number(), cve.Metadata.Assigner)
  if  err != nil {
    return err
  }

  // get cve ID
  cveId, err = rs.LastInsertId()
  if err != nil {
    return err
  }

  // add problem types
  for _, pt := range(cve.ProblemTypes.ProblemTypes) {
    // problem type ID
    var ptId int64

    // add problem type
    rs, err := tx.Exec(ctx, "feed/insert-cve-problem", cveId)
    if err != nil {
      return err
    }

    // get problem type ID
    ptId, err = rs.LastInsertId()
    if err != nil {
      return err
    }

    // add problem type descriptions
    descIds, err := me.addDescriptions(ctx, tx, pt.Descriptions)
    if err != nil {
      return err
    }

    // add problem type descriptions
    for _, descId := range(descIds) {
      _, err = tx.Exec(ctx, "feed/insert-cve-problem-desc", ptId, descId)
      if err != nil {
        return err
      }
    }
  }

  // add references
  for _, ref := range(cve.References.References) {
    // add reference
    rs, err := tx.Exec(ctx, "feed/insert-cve-ref", cveId, ref.Url, ref.Name, ref.RefSource)
    if err != nil {
      return err
    }

    // get ref ID
    refId, err := rs.LastInsertId()
    if err != nil {
      return err
    }

    // add cve ref tags
    for _, tag := range(ref.Tags) {
      _, err = tx.Exec(ctx, "feed/insert-cve-ref-tag", refId, tag)
      if err != nil {
        return err
      }
    }
  }

  // add cve descriptions
  descIds, err := me.addDescriptions(ctx, tx, cve.Description.Descriptions)
  if err != nil {
    return err
  }

  // add cve descriptions
  for _, descId := range(descIds) {
    _, err = tx.Exec(ctx, "feed/insert-cve-desc", cveId, descId)
    if err != nil {
      return err
    }
  }

  // return success
  return nil
}

// add feed item CVSSv2 impact.
func (me DbStore) addItemCvssV2(ctx context.Context, tx Tx, itemId int64, metric *nvd_feed.BaseMetricV2) error {
  if metric == nil {
    // no v2 vector to add
    return nil
  }

  // check version
  if metric.CvssV2.Version != nvd_feed.V20 {
    return fmt.Errorf("unknown CVSSv2 version: %s", metric.CvssV2.Version)
  }

  // add metric
  _, err := tx.Exec(ctx, "feed/insert-item-cvss-v2",
    itemId,
    metric.CvssV2.Vector.String(),
    int64(metric.CvssV2.BaseScore),
    metric.Severity.String(),
    int64(metric.ExploitabilityScore),
    int64(metric.ImpactScore),
    metric.InsufficientInfo,
    metric.ObtainAllPrivilege,
    metric.ObtainUserPrivilege,
    metric.ObtainOtherPrivilege,
    metric.UserInteractionRequired,
  )

  // return result
  return err
}

// add feed item CVSSv3 impact.
func (me DbStore) addItemCvssV3(ctx context.Context, tx Tx, itemId int64, metric *nvd_feed.BaseMetricV3) error {
  if metric == nil {
    // no cvssv3 vector to add
    return nil
  }

  // check version
  if metric.CvssV3.Version != nvd_feed.V30 && metric.CvssV3.Version != nvd_feed.V31 {
    return fmt.Errorf("unknown CVSSv3 version: %s", metric.CvssV3.Version)
  }

  // add metric
  _, err := tx.Exec(ctx, "feed/insert-item-cvss-v3",
    itemId,
    metric.CvssV3.Vector.String(),
    int64(uint8(metric.CvssV3.BaseScore)),
    metric.CvssV3.BaseSeverity.String(),
    int64(uint8(metric.ExploitabilityScore)),
    int64(uint8(metric.ImpactScore)),
  )

  // return result
  return err
}

// add feed item impact.
func (me DbStore) addImpact(ctx context.Context, tx Tx, itemId int64, impact nvd_feed.Impact) error {
  // add cvss v2 impact
  if err := me.addItemCvssV2(ctx, tx, itemId, impact.BaseMetricV2); err != nil {
    return err
  }

  // add cvss v3 impact
  if err := me.addItemCvssV3(ctx, tx, itemId, impact.BaseMetricV3); err != nil {
    return err
  }

  // return success
  return nil
}

// Add feed item and return item ID.
func (me DbStore) addItem(ctx context.Context, tx Tx, feedId int64, item nvd_feed.Item) error {
  // item ID
  var itemId int64

  // serialize published timestamp
  publishedAt, err := item.PublishedDate.MarshalText()
  if err != nil {
    return err
  }

  // serialize lastmodified timestamp
  modifiedAt, err := item.LastModifiedDate.MarshalText()
  if err != nil {
    return err
  }

  // add feed item
  rs, err := tx.Exec(ctx, "feed/insert-item", feedId, publishedAt, modifiedAt)
  if  err != nil {
    return err
  }

  // get item ID
  itemId, err = rs.LastInsertId()
  if err != nil {
    return err
  }

  // TODO: add item configuration
  // if err := me.addConfigs(ctx, tx, itemId, item.Configurations); err != nil {
  //   return err
  // }


  // add item cve
  if err := me.addCve(ctx, tx, itemId, item.Cve); err != nil {
    return err
  }

  // add item impact
  if err := me.addImpact(ctx, tx, itemId, item.Impact); err != nil {
    return err
  }

  // return success
  return nil
}

// Add base CVE feed and return feed ID.
func (me DbStore) addFeed(ctx context.Context, tx Tx, feed nvd_feed.Feed) (int64, error) {
  // return feed ID
  var feedId int64

  // check feed data type, data format, and data version
  if err := checkNvdData(feed.DataType, feed.DataFormat, feed.DataVersion); err != nil {
    return feedId, err
  }

  // FIXME: check cve count?
  // if feed.NumCVEs != len(feed.Items) {
  // }

  // serialize feed timestamp
  ts, err := feed.Timestamp.MarshalText()
  if err != nil {
    return feedId, err
  }

  // add feed
  rs, err := tx.Exec(ctx, "feed/insert", string(ts))
  if err != nil {
    return feedId, err
  }

  // get feed ID
  feedId, err = rs.LastInsertId()
  if err != nil {
    return feedId, err
  }

  // add feed items
  for _, item := range(feed.Items) {
    if err := me.addItem(ctx, tx, feedId, item); err != nil {
      return feedId, err
    }
  }

  // return success
  return feedId, nil
}

// query IDs used by AddCpeMatches()
var addCveFeedQueryIds = []string {
  "feed/insert",
  "feed/insert-assigner",
  "feed/insert-cve",
  "feed/insert-cve-desc",
  "feed/insert-cve-fts-refresh",
  "feed/insert-cve-problem",
  "feed/insert-cve-problem-desc",
  "feed/insert-cve-ref",
  "feed/insert-cve-ref-tag",
  "feed/insert-desc",
  "feed/insert-item",
  "feed/insert-item-cvss-v2",
  "feed/insert-item-cvss-v3",
}

// Refresh CVE FTS index.
func (me DbStore) refreshCveFts(ctx context.Context, tx Tx) error {
  _, err := tx.Exec(ctx, "feed/insert-cve-fts-refresh")
  return err
}

// Import CVE feeds.
func (me DbStore) AddCveFeeds(ctx context.Context, feeds []nvd_feed.Feed) ([]int64, error) {
  feedIds := make([]int64, len(feeds))

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

  // begin transaction
  err := me.Tx(ctx, addCveFeedQueryIds, func(tx Tx) error {
    for i, feed := range(feeds) {
      // add feed, get feed ID
      if id, err := me.addFeed(ctx, tx, feed); err != nil {
        return err
      } else {
        feedIds[i] = id
      }
    }

    // refresh the cve fts index
    return me.refreshCveFts(ctx, tx)
  })

  // return results
  return feedIds, err
}

// Import single CVE feed.
func (me DbStore) AddCveFeed(ctx context.Context, feed nvd_feed.Feed) (int64, error) {
  if r, err := me.AddCveFeeds(ctx, []nvd_feed.Feed { feed }); err != nil {
    return 0, err
  } else {
    return r[0], nil
  }
}

// search CVEs
func (me DbStore) CveSearch(
  ctx context.Context,
  s string,
) ([]CveSearchRow, error) {
  var r []CveSearchRow

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

  // get/exec search query
  err := me.QueryRow(ctx, "feed/search", []interface{} {
    db_sql.Named("q", s),
  }, func(row *db_sql.Row) error {
    var s string
    if err := row.Scan(&s); err != nil {
      return err
    }

    return json.Unmarshal([]byte(s), &r)
  })

  // return results
  return r, err
}