Merge pull request #61 from bitvector2/master

added saving results to PostgreSQL DB as a JSONB document

cmd/common.go

@@ -103,10 +103,20 @@ func runChecks(t check.NodeType) {
 		}
 
 		fmt.Println(string(out))
+	} else {
+		// if we want to store in PostgreSQL, convert to JSON and save it
+		if (summary.Fail > 0 || summary.Warn > 0 || summary.Pass > 0) && pgSql {
+			out, err := controls.JSON()
+			if err != nil {
+				exitWithError(fmt.Errorf("failed to output in JSON format: %v", err))
+			}
+
+			savePgsql(string(out))
 		} else {
 			prettyPrint(controls, summary)
 		}
 	}
+}
 
 // colorPrint outputs the state in a specific colour, along with a message string
 func colorPrint(state check.State, s string) {

cmd/database.go

@@ -0,0 +1,60 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/golang/glog"
+	"github.com/jinzhu/gorm"
+	_ "github.com/jinzhu/gorm/dialects/postgres"
+	"github.com/spf13/viper"
+)
+
+func savePgsql(jsonInfo string) {
+	envVars := map[string]string{
+		"PGSQL_HOST":     viper.GetString("PGSQL_HOST"),
+		"PGSQL_USER":     viper.GetString("PGSQL_USER"),
+		"PGSQL_DBNAME":   viper.GetString("PGSQL_DBNAME"),
+		"PGSQL_SSLMODE":  viper.GetString("PGSQL_SSLMODE"),
+		"PGSQL_PASSWORD": viper.GetString("PGSQL_PASSWORD"),
+	}
+
+	for k, v := range envVars {
+		if v == "" {
+			exitWithError(fmt.Errorf("environment variable %s is missing", envVarsPrefix+"_"+k))
+		}
+	}
+
+	connInfo := fmt.Sprintf("host=%s user=%s dbname=%s sslmode=%s password=%s",
+		envVars["PGSQL_HOST"],
+		envVars["PGSQL_USER"],
+		envVars["PGSQL_DBNAME"],
+		envVars["PGSQL_SSLMODE"],
+		envVars["PGSQL_PASSWORD"],
+	)
+
+	hostname, err := os.Hostname()
+	if err != nil {
+		exitWithError(fmt.Errorf("received error looking up hostname: %s", err))
+	}
+
+	timestamp := time.Now()
+
+	type ScanResult struct {
+		gorm.Model
+		ScanHost string    `gorm:"type:varchar(63) not null"` // https://www.ietf.org/rfc/rfc1035.txt
+		ScanTime time.Time `gorm:"not null"`
+		ScanInfo string    `gorm:"type:jsonb not null"`
+	}
+
+	db, err := gorm.Open("postgres", connInfo)
+	defer db.Close()
+	if err != nil {
+		exitWithError(fmt.Errorf("received error connecting to database: %s", err))
+	}
+
+	db.Debug().AutoMigrate(&ScanResult{})
+	db.Save(&ScanResult{ScanHost: hostname, ScanTime: timestamp, ScanInfo: jsonInfo})
+	glog.V(2).Info(fmt.Sprintf("successfully stored result to: %s", envVars["PGSQL_HOST"]))
+}

cmd/root.go

@@ -25,10 +25,11 @@ import (
 )
 
 var (
+	envVarsPrefix = "KUBE_BENCH"
 	cfgDir        = "./cfg"
 	cfgFile       string
-
 	jsonFmt       bool
+	pgSql         bool
 	checkList     string
 	groupList     string
 	masterFile    string
@@ -59,6 +60,7 @@ func init() {
 	cobra.OnInitialize(initConfig)
 
 	RootCmd.PersistentFlags().BoolVar(&jsonFmt, "json", false, "Prints the results as JSON")
+	RootCmd.PersistentFlags().BoolVar(&pgSql, "pgsql", false, "Save the results to PostgreSQL")
 	RootCmd.PersistentFlags().StringVarP(
 		&checkList,
 		"check",
@@ -90,7 +92,7 @@ func initConfig() {
 		viper.AddConfigPath(cfgDir)   // adding ./cfg as first search path
 	}
 
-	viper.SetEnvPrefix("KUBE_BENCH")
+	viper.SetEnvPrefix(envVarsPrefix)
 	viper.AutomaticEnv() // read in environment variables that match
 
 	// If a config file is found, read it in.