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  }