<commit_before>package db

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	log "github.com/kotakanbe/go-cve-dictionary/log"

	"github.com/jinzhu/gorm"
	"github.com/k0kubun/pp"
	c "github.com/kotakanbe/go-cve-dictionary/config"
	jvn "github.com/kotakanbe/go-cve-dictionary/jvn"
	"github.com/kotakanbe/go-cve-dictionary/models"
	"github.com/kotakanbe/go-cve-dictionary/nvd"

	// Required MySQL.  See http://jinzhu.me/gorm/database.html#connecting-to-a-database
	_ "github.com/jinzhu/gorm/dialects/mysql"
	// Required SQLite3.
	_ "github.com/jinzhu/gorm/dialects/sqlite"
)

var db *gorm.DB

// Supported DB dialects.
const (
	dialectSqlite3 = "sqlite3"
	dialectMysql   = "mysql"
)

// OpenDB opens Database
func OpenDB() (err error) {

	db, err = gorm.Open(c.Conf.DBType, c.Conf.DBPath)
	if err != nil {
		if c.Conf.DBType == dialectSqlite3 {
			err = fmt.Errorf("Failed to open DB. datafile: %s, err: %s", c.Conf.DBPath, err)
		} else if c.Conf.DBType == dialectMysql {
			err = fmt.Errorf("Failed to open DB, err: %s", err)
		} else {
			err = fmt.Errorf("Invalid database dialect, %s", c.Conf.DBType)
		}
		return
	}

	db.LogMode(c.Conf.DebugSQL)

	if c.Conf.DBType == dialectSqlite3 {
		db.Exec("PRAGMA journal_mode=WAL;")
	}

	return
}

func recconectDB() error {
	var err error
	if err = db.Close(); err != nil {
		return fmt.Errorf("Failed to close DB. Type: %s, Path: %s, err: %s", c.Conf.DBType, c.Conf.DBPath, err)
	}
	return OpenDB()
}

// MigrateDB migrates Database
func MigrateDB() error {
	if err := db.AutoMigrate(
		&models.CveDetail{},
		&models.Jvn{},
		&models.Nvd{},
		&models.Reference{},
		&models.Cpe{},
	).Error; err != nil {
		return fmt.Errorf("Failed to migrate. err: %s", err)
	}

	errMsg := "Failed to create index. err: %s"
	if err := db.Model(&models.CveDetail{}).
		AddIndex("idx_cve_detail_cveid", "cve_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Nvd{}).
		AddIndex("idx_nvds_cve_detail_id", "cve_detail_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Jvn{}).
		AddIndex("idx_jvns_cve_detail_id", "cve_detail_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Cpe{}).
		AddIndex("idx_cpes_jvn_id", "jvn_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Reference{}).
		AddIndex("idx_references_jvn_id", "jvn_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Cpe{}).
		AddIndex("idx_cpes_nvd_id", "nvd_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Cpe{}).
		AddIndex("idx_cpes_cpe_name", "cpe_name").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Reference{}).
		AddIndex("idx_references_nvd_id", "nvd_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}

	return nil
}

// Get Select Cve information from DB.
func Get(cveID string, priorityDB ...*gorm.DB) models.CveDetail {

	var conn *gorm.DB
	if len(priorityDB) == 1 {
		conn = priorityDB[0]
	} else {
		conn = db
	}

	c := models.CveDetail{}
	conn.Where(&models.CveDetail{CveID: cveID}).First(&c)

	if c.ID == 0 {
		return models.CveDetail{}
	}

	// JVN
	jvn := models.Jvn{}
	conn.Model(&c).Related(&jvn, "Jvn")
	c.Jvn = jvn

	if jvn.CveDetailID != 0 && jvn.ID != 0 {
		jvnRefs := []models.Reference{}
		conn.Model(&jvn).Related(&jvnRefs, "References")
		c.Jvn.References = jvnRefs

		// TODO commentout because JSON response size will be big. so Uncomment if needed.
		//  jvnCpes := []models.Cpe{}
		//  conn.Model(&jvn).Related(&jvnCpes, "Cpes")
		//  c.Jvn.Cpes = jvnCpes
	}

	// NVD
	nvd := models.Nvd{}
	conn.Model(&c).Related(&nvd, "Nvd")
	c.Nvd = nvd

	if nvd.CveDetailID != 0 && nvd.ID != 0 {
		nvdRefs := []models.Reference{}
		conn.Model(&nvd).Related(&nvdRefs, "References")
		c.Nvd.References = nvdRefs

		// TODO commentout because JSON response size will be big. so Uncomment if needed.
		//  nvdCpes := []models.Cpe{}
		//  conn.Model(&nvd).Related(&nvdCpes, "Cpes")
		//  c.Nvd.Cpes = nvdCpes
	}

	return c
}

// GetByCpeName Select Cve information from DB.
func GetByCpeName(cpeName string) (details []models.CveDetail) {
	cpes := []models.Cpe{}
	db.Where(&models.Cpe{CpeName: cpeName}).Find(&cpes)

	for _, cpe := range cpes {
		var cveDetailID uint
		if cpe.JvnID != 0 {
			//TODO test check CPE name format of JVN table
			jvn := models.Jvn{}
			db.Select("cve_detail_id").Where("ID = ?", cpe.JvnID).First(&jvn)
			cveDetailID = jvn.CveDetailID
		} else if cpe.NvdID != 0 {
			nvd := models.Nvd{}
			db.Select("cve_detail_id").Where("ID = ?", cpe.NvdID).First(&nvd)
			cveDetailID = nvd.CveDetailID
		}<commit_msg>Avoid null slice being null in JSON<commit_after>package db

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	log "github.com/kotakanbe/go-cve-dictionary/log"

	"github.com/jinzhu/gorm"
	"github.com/k0kubun/pp"
	c "github.com/kotakanbe/go-cve-dictionary/config"
	jvn "github.com/kotakanbe/go-cve-dictionary/jvn"
	"github.com/kotakanbe/go-cve-dictionary/models"
	"github.com/kotakanbe/go-cve-dictionary/nvd"

	// Required MySQL.  See http://jinzhu.me/gorm/database.html#connecting-to-a-database
	_ "github.com/jinzhu/gorm/dialects/mysql"
	// Required SQLite3.
	_ "github.com/jinzhu/gorm/dialects/sqlite"
)

var db *gorm.DB

// Supported DB dialects.
const (
	dialectSqlite3 = "sqlite3"
	dialectMysql   = "mysql"
)

// OpenDB opens Database
func OpenDB() (err error) {

	db, err = gorm.Open(c.Conf.DBType, c.Conf.DBPath)
	if err != nil {
		if c.Conf.DBType == dialectSqlite3 {
			err = fmt.Errorf("Failed to open DB. datafile: %s, err: %s", c.Conf.DBPath, err)
		} else if c.Conf.DBType == dialectMysql {
			err = fmt.Errorf("Failed to open DB, err: %s", err)
		} else {
			err = fmt.Errorf("Invalid database dialect, %s", c.Conf.DBType)
		}
		return
	}

	db.LogMode(c.Conf.DebugSQL)

	if c.Conf.DBType == dialectSqlite3 {
		db.Exec("PRAGMA journal_mode=WAL;")
	}

	return
}

func recconectDB() error {
	var err error
	if err = db.Close(); err != nil {
		return fmt.Errorf("Failed to close DB. Type: %s, Path: %s, err: %s", c.Conf.DBType, c.Conf.DBPath, err)
	}
	return OpenDB()
}

// MigrateDB migrates Database
func MigrateDB() error {
	if err := db.AutoMigrate(
		&models.CveDetail{},
		&models.Jvn{},
		&models.Nvd{},
		&models.Reference{},
		&models.Cpe{},
	).Error; err != nil {
		return fmt.Errorf("Failed to migrate. err: %s", err)
	}

	errMsg := "Failed to create index. err: %s"
	if err := db.Model(&models.CveDetail{}).
		AddIndex("idx_cve_detail_cveid", "cve_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Nvd{}).
		AddIndex("idx_nvds_cve_detail_id", "cve_detail_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Jvn{}).
		AddIndex("idx_jvns_cve_detail_id", "cve_detail_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Cpe{}).
		AddIndex("idx_cpes_jvn_id", "jvn_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Reference{}).
		AddIndex("idx_references_jvn_id", "jvn_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Cpe{}).
		AddIndex("idx_cpes_nvd_id", "nvd_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Cpe{}).
		AddIndex("idx_cpes_cpe_name", "cpe_name").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}
	if err := db.Model(&models.Reference{}).
		AddIndex("idx_references_nvd_id", "nvd_id").Error; err != nil {
		return fmt.Errorf(errMsg, err)
	}

	return nil
}

// Get Select Cve information from DB.
func Get(cveID string, priorityDB ...*gorm.DB) models.CveDetail {
	var conn *gorm.DB
	if len(priorityDB) == 1 {
		conn = priorityDB[0]
	} else {
		conn = db
	}

	c := models.CveDetail{}
	conn.Where(&models.CveDetail{CveID: cveID}).First(&c)

	if c.ID == 0 {
		// Avoid null slice being null in JSON
		return models.CveDetail{
			Nvd: models.Nvd{
				References: []models.Reference{},
				Cpes:       []models.Cpe{},
			},
			Jvn: models.Jvn{
				References: []models.Reference{},
				Cpes:       []models.Cpe{},
			},
		}
	}

	// JVN
	jvn := models.Jvn{}
	conn.Model(&c).Related(&jvn, "Jvn")
	c.Jvn = jvn

	if jvn.CveDetailID != 0 && jvn.ID != 0 {
		jvnRefs := []models.Reference{}
		conn.Model(&jvn).Related(&jvnRefs, "References")
		c.Jvn.References = jvnRefs

		// TODO commentout because JSON response size will be big. so Uncomment if needed.
		//  jvnCpes := []models.Cpe{}
		//  conn.Model(&jvn).Related(&jvnCpes, "Cpes")
		//  c.Jvn.Cpes = jvnCpes
		if c.Jvn.Cpes == nil {
			c.Jvn.Cpes = []models.Cpe{}
		}
	}

	// NVD
	nvd := models.Nvd{}
	conn.Model(&c).Related(&nvd, "Nvd")
	c.Nvd = nvd

	if nvd.CveDetailID != 0 && nvd.ID != 0 {
		nvdRefs := []models.Reference{}
		conn.Model(&nvd).Related(&nvdRefs, "References")
		c.Nvd.References = nvdRefs

		// TODO commentout because JSON response size will be big. so Uncomment if needed.
		//  nvdCpes := []models.Cpe{}
		//  conn.Model(&nvd).Related(&nvdCpes, "Cpes")
		//  c.Nvd.Cpes = nvdCpes
		if c.Nvd.Cpes == nil {
			c.Nvd.Cpes = []models.Cpe{}
		}
	}

	return c
}

// GetByCpeName Select Cve information from DB.
func GetByCpeName(cpeName string) (details []models.CveDetail) {
	cpes := []models.Cpe{}
	db.Where(&models.Cpe{CpeName: cpeName}).Find(&cpes)

	for _, cpe := range cpes {
		var cveDetailID uint
		if cpe.JvnID != 0 {
			//TODO test check CPE name format of JVN table
			jvn := models.Jvn{}
			db.Select("cve_detail_id").Where("ID = ?", cpe.JvnID).First(&jvn)
			cveDetailID = jvn.CveDetailID
		} else if cpe.NvdID != 0 {
			nvd := models.Nvd{}
			db.Select("cve_detail_id").Where("ID = ?", cpe.NvdID).First(&nvd)
			cveDetailID = nvd.CveDetailID
		}
{'fraction_non_alphanumeric': 0.17463701898669914, 'fraction_numerical': 0.003045994517209869, 'mean_word_length': 3.5496535796766744, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5666925', 'n_tokens_mistral': 4426, 'n_tokens_neox': 4063, 'n_words': 815}