db.go

package main

import (
	"encoding/xml"
	"fmt"
	"os"

	"github.com/jinzhu/gorm"
	"github.com/jordan-wright/cve-api/cve"
	_ "github.com/mattn/go-sqlite3"
)

var db gorm.DB

func init() {
	var err error
	db, err = gorm.Open("sqlite3", "cve.db")
	if err != nil {
		panic(fmt.Sprintf("Got error when connect database, the error is '%v'", err))
	}
	if _, err = os.Stat("cve.db"); err != nil {
		fmt.Println("Creating database..")
		db.CreateTable(cve.CVSS{})
		db.CreateTable(cve.Reference{})
		db.CreateTable(cve.Link{})
		db.CreateTable(cve.Product{})
		db.CreateTable(cve.Entry{})
		fmt.Println("Unmarshalling XML entries into database")
		tx := db.Begin()
		f, err := os.OpenFile("cve/nvdcve-2.0-2002.xml", os.O_RDONLY, 0)
		if err != nil {
			fmt.Println(err)
		}
		decoder := xml.NewDecoder(f)
		for {
			// Read tokens from the XML document in a stream.
			t, _ := decoder.Token()
			if t == nil {
				break
			}
			// Inspect the type of the token just read.
			switch se := t.(type) {
			case xml.StartElement:
				if se.Name.Local == "entry" {
					var e cve.Entry
					// decode a whole chunk of following XML into the
					// entry
					err = decoder.DecodeElement(&e, &se)
					if err != nil {
						fmt.Println(err)
					}
					tx.Save(&e)
					// Save all the products (needs to be done separately for now)
					for _, p := range e.Products {
						tx.Save(&cve.Product{
							EntryId: e.Id,
							Value:   p,
						})
					}
				}
			}
		}
		err = tx.Commit().Error
		if err != nil {
			fmt.Println(err)
		}
	}
}