github.com/khulnasoft-lab/kube-bench@v0.2.1-0.20240330183753-9df52345ae58/cmd/database.go (about) 1 package cmd 2 3 import ( 4 "fmt" 5 "os" 6 "time" 7 8 "github.com/golang/glog" 9 "github.com/spf13/viper" 10 "gorm.io/driver/postgres" 11 "gorm.io/gorm" 12 ) 13 14 type PsqlConnInfo struct { 15 Host string 16 User string 17 DbName string 18 SslMode string 19 Password string 20 } 21 22 func getPsqlConnInfo() (PsqlConnInfo, error) { 23 var host string 24 if value := viper.GetString("PGSQL_HOST"); value != "" { 25 host = value 26 } else { 27 return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_HOST env var is required", envVarsPrefix) 28 } 29 30 var user string 31 if value := viper.GetString("PGSQL_USER"); value != "" { 32 user = value 33 } else { 34 return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_USER env var is required", envVarsPrefix) 35 } 36 37 var dbName string 38 if value := viper.GetString("PGSQL_DBNAME"); value != "" { 39 dbName = value 40 } else { 41 return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_DBNAME env var is required", envVarsPrefix) 42 } 43 44 var sslMode string 45 if value := viper.GetString("PGSQL_SSLMODE"); value != "" { 46 sslMode = value 47 } else { 48 return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_SSLMODE env var is required", envVarsPrefix) 49 } 50 51 var password string 52 if value := viper.GetString("PGSQL_PASSWORD"); value != "" { 53 password = value 54 } else { 55 return PsqlConnInfo{}, fmt.Errorf("%s_PGSQL_PASSWORD env var is required", envVarsPrefix) 56 } 57 58 return PsqlConnInfo{ 59 Host: host, 60 User: user, 61 DbName: dbName, 62 SslMode: sslMode, 63 Password: password, 64 }, nil 65 } 66 67 func (c *PsqlConnInfo) toString() string { 68 return fmt.Sprintf("host=%s user=%s dbname=%s sslmode=%s password=%s", 69 c.Host, 70 c.User, 71 c.DbName, 72 c.SslMode, 73 c.Password, 74 ) 75 } 76 77 func savePgsql(jsonInfo string) { 78 var hostname string 79 if value := viper.GetString("K8S_HOST"); value != "" { 80 // Adhere to the ScanHost column definition below 81 if len(value) > 63 { 82 exitWithError(fmt.Errorf("%s_K8S_HOST value's length must be less than 63 chars", envVarsPrefix)) 83 } 84 85 hostname = value 86 } else { 87 host, err := os.Hostname() 88 if err != nil { 89 exitWithError(fmt.Errorf("received error looking up hostname: %s", err)) 90 } 91 92 hostname = host 93 } 94 95 PsqlConnInfo, err := getPsqlConnInfo() 96 if err != nil { 97 exitWithError(err) 98 } 99 100 db, err := gorm.Open(postgres.Open(PsqlConnInfo.toString()), &gorm.Config{}) 101 if err != nil { 102 exitWithError(fmt.Errorf("received error connecting to database: %s", err)) 103 } 104 105 timestamp := time.Now() 106 type ScanResult struct { 107 gorm.Model 108 ScanHost string `gorm:"type:varchar(63) not null"` // https://www.ietf.org/rfc/rfc1035.txt 109 ScanTime time.Time `gorm:"not null"` 110 ScanInfo string `gorm:"type:jsonb not null"` 111 } 112 113 db.Debug().AutoMigrate(&ScanResult{}) 114 db.Save(&ScanResult{ScanHost: hostname, ScanTime: timestamp, ScanInfo: jsonInfo}) 115 glog.V(2).Info(fmt.Sprintf("successfully stored result to: %s", PsqlConnInfo.Host)) 116 }