Read sftp key/keyfile and expose via config

cmd/root.go

@@ -2,13 +2,15 @@ package cmd
 
 import (
 	"fmt"
+	"strings"
+
 	"github.com/arunvelsriram/sftp-exporter/pkg/config"
 	c "github.com/arunvelsriram/sftp-exporter/pkg/constants"
 	"github.com/arunvelsriram/sftp-exporter/pkg/server"
 	log "github.com/sirupsen/logrus"
+	"github.com/spf13/afero"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-	"strings"
 )
 
 var cfg config.Config
@@ -59,13 +61,29 @@ func init() {
 	rootCmd.PersistentFlags().String(c.FlagSFTPUser, "", "sftp user")
 	_ = viper.BindPFlag(c.ViperKeySFTPUser, rootCmd.PersistentFlags().Lookup(c.FlagSFTPUser))
 
-	rootCmd.PersistentFlags().String(c.FlagSFTPPass, "", "sftp user")
+	rootCmd.PersistentFlags().String(c.FlagSFTPPass, "", "sftp password")
 	_ = viper.BindPFlag(c.ViperKeySFTPPass, rootCmd.PersistentFlags().Lookup(c.FlagSFTPPass))
+
+	rootCmd.PersistentFlags().String(c.FlagSFTPKey, "", "sftp key (base64 encoded)")
+	_ = viper.BindPFlag(c.ViperKeySFTPKey, rootCmd.PersistentFlags().Lookup(c.FlagSFTPKey))
+
+	rootCmd.PersistentFlags().String(c.FlagSFTPKeyFile, "", "sftp key file")
+	_ = viper.BindPFlag(c.ViperKeySFTPKeyFile, rootCmd.PersistentFlags().Lookup(c.FlagSFTPKeyFile))
+
+	rootCmd.PersistentFlags().String(c.FlagSFTPKeyPassphrase, "", "sftp key passphrase")
+	_ = viper.BindPFlag(c.ViperKeySFTPKeyPassphrase, rootCmd.PersistentFlags().Lookup(c.FlagSFTPKeyPassphrase))
 }
 
 func initConfig() {
 	viper.AutomaticEnv()
-	cfg = config.LoadConfig()
+
+	var err error
+	fs := afero.NewOsFs()
+	cfg, err = config.LoadConfig(fs)
+	if err != nil {
+		panic(err)
+	}
+
 	log.SetFormatter(&log.TextFormatter{
 		FullTimestamp: true,
 	})

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/pkg/sftp v1.11.0
 	github.com/prometheus/client_golang v0.9.3
 	github.com/sirupsen/logrus v1.2.0
+	github.com/spf13/afero v1.1.2
 	github.com/spf13/cobra v1.0.0
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/spf13/viper v1.4.0

pkg/client/sftp_client.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"fmt"
+
 	log "github.com/sirupsen/logrus"
 
 	"github.com/arunvelsriram/sftp-exporter/pkg/config"

pkg/config/config.go

@@ -1,15 +1,23 @@
 package config
 
 import (
+	"encoding/base64"
+	"fmt"
+
 	c "github.com/arunvelsriram/sftp-exporter/pkg/constants"
+	"github.com/arunvelsriram/sftp-exporter/pkg/utils"
+	"github.com/spf13/afero"
 	"github.com/spf13/viper"
 )
 
 type sftpConfig struct {
-	Host string
-	Port int
-	User string
-	Pass string
+	Host          string
+	Port          int
+	User          string
+	Pass          string
+	Key           []byte
+	KeyFile       string
+	KeyPassphrase string
 }
 
 type Config interface {
@@ -20,6 +28,9 @@ type Config interface {
 	GetSFTPPort() int
 	GetSFTPUser() string
 	GetSFTPPass() string
+	GetSFTPKey() []byte
+	GetSFTPKeyFile() string
+	GetSFTPKeyPassphrase() string
 }
 
 type sftpExporterConfig struct {
@@ -29,18 +40,47 @@ type sftpExporterConfig struct {
 	SFTPConfig  sftpConfig
 }
 
-func LoadConfig() Config {
+func resolveKey(encodedKey, keyfile string, fs afero.Fs) (sftpKey []byte, err error) {
+	if utils.IsNotEmpty(encodedKey) && utils.IsNotEmpty(keyfile) {
+		return sftpKey, fmt.Errorf("only one of key or keyfile should be specified")
+	}
+
+	if utils.IsNotEmpty(encodedKey) {
+		sftpKey, err = base64.StdEncoding.DecodeString(encodedKey)
+		if err != nil {
+			return sftpKey, err
+		}
+	} else if utils.IsNotEmpty(keyfile) {
+		sftpKey, err = afero.ReadFile(fs, keyfile)
+		if err != nil {
+			return sftpKey, err
+		}
+	}
+	return sftpKey, nil
+}
+
+func LoadConfig(fs afero.Fs) (Config, error) {
+	encodedKey := viper.GetString(c.ViperKeySFTPKey)
+	keyFile := viper.GetString(c.ViperKeySFTPKeyFile)
+	key, err := resolveKey(encodedKey, keyFile, fs)
+	if err != nil {
+		return nil, err
+	}
+
 	return sftpExporterConfig{
 		BindAddress: viper.GetString(c.ViperKeyBindAddress),
 		Port:        viper.GetInt(c.ViperKeyPort),
 		LogLevel:    viper.GetString(c.ViperKeyLogLevel),
 		SFTPConfig: sftpConfig{
-			Host: viper.GetString(c.ViperKeySFTPHost),
-			Port: viper.GetInt(c.ViperKeySFTPPort),
-			User: viper.GetString(c.ViperKeySFTPUser),
-			Pass: viper.GetString(c.ViperKeySFTPPass),
+			Host:          viper.GetString(c.ViperKeySFTPHost),
+			Port:          viper.GetInt(c.ViperKeySFTPPort),
+			User:          viper.GetString(c.ViperKeySFTPUser),
+			Pass:          viper.GetString(c.ViperKeySFTPPass),
+			Key:           key,
+			KeyFile:       keyFile,
+			KeyPassphrase: viper.GetString(c.ViperKeySFTPKeyPassphrase),
 		},
-	}
+	}, nil
 }
 
 func (c sftpExporterConfig) GetBindAddress() string {
@@ -69,4 +109,16 @@ func (c sftpExporterConfig) GetSFTPUser() string {
 
 func (c sftpExporterConfig) GetSFTPPass() string {
 	return c.SFTPConfig.Pass
-}
+}
+
+func (c sftpExporterConfig) GetSFTPKey() []byte {
+	return c.SFTPConfig.Key
+}
+
+func (c sftpExporterConfig) GetSFTPKeyFile() string {
+	return c.SFTPConfig.KeyFile
+}
+
+func (c sftpExporterConfig) GetSFTPKeyPassphrase() string {
+	return c.SFTPConfig.KeyPassphrase
+}

pkg/config/config_test.go

@@ -3,30 +3,81 @@ package config
 import (
 	"testing"
 
-	. "github.com/arunvelsriram/sftp-exporter/pkg/constants"
+	c "github.com/arunvelsriram/sftp-exporter/pkg/constants"
+	"github.com/spf13/afero"
 	"github.com/spf13/viper"
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/suite"
 )
 
-func TestNewConfig(t *testing.T) {
-	viper.Set(ViperKeyBindAddress, "127.0.0.1")
-	viper.Set(ViperKeyPort, 8080)
-	viper.Set(ViperKeySFTPHost, "localhost")
-	viper.Set(ViperKeySFTPPort, 22)
-	viper.Set(ViperKeySFTPUser, "arun")
-	viper.Set(ViperKeySFTPPass, "arun@123")
-
-	c := LoadConfig()
-
-	expected := sftpExporterConfig{
-		BindAddress: "127.0.0.1",
-		Port:        8080,
-		SFTPConfig: sftpConfig{
-			Host: "localhost",
-			Port: 22,
-			User: "arun",
-			Pass: "arun@123",
-		},
-	}
-	assert.Equal(t, expected, c)
+type ConfigTestSuite struct {
+	suite.Suite
+	fs afero.Fs
+}
+
+func TestConfigTestSuite(t *testing.T) {
+	suite.Run(t, new(ConfigTestSuite))
+}
+
+func (s *ConfigTestSuite) SetupTest() {
+	s.fs = afero.NewMemMapFs()
+}
+
+func (s *ConfigTestSuite) TestConfigLoadConfig() {
+	s.Run("should return config", func() {
+		actual, err := LoadConfig(s.fs)
+
+		s.NoError(err)
+		s.Equal(sftpExporterConfig{}, actual)
+	})
+
+	s.Run("should return err when both key and keyfile are provided", func() {
+		viper.Reset()
+		viper.Set(c.ViperKeySFTPKey, "sftp private key")
+		viper.Set(c.ViperKeySFTPKeyFile, "sftp private keyfile")
+
+		_, err := LoadConfig(s.fs)
+
+		s.EqualError(err, "only one of key or keyfile should be specified")
+	})
+
+	s.Run("should return error if key is not encoded properly", func() {
+		viper.Reset()
+		viper.Set(c.ViperKeySFTPKey, "invalid encoding")
+
+		_, err := LoadConfig(s.fs)
+
+		s.EqualError(err, "illegal base64 data at input byte 7")
+	})
+
+	s.Run("should store decoded key for given encoded key", func() {
+		viper.Reset()
+		viper.Set(c.ViperKeySFTPKey, "YXJ1bg==")
+
+		actual, err := LoadConfig(s.fs)
+
+		s.NoError(err)
+		s.Equal([]byte("arun"), actual.GetSFTPKey())
+	})
+
+	s.Run("should return error if reading keyfile fails", func() {
+		viper.Reset()
+		viper.Set(c.ViperKeySFTPKeyFile, "invalidfile")
+
+		_, err := LoadConfig(s.fs)
+
+		s.EqualError(err, "open invalidfile: file does not exist")
+	})
+
+	s.Run("should store key and keyfile for a valid keyfile", func() {
+		file, _ := afero.TempFile(s.fs, "", "config-test")
+		_, _ = file.WriteString("private key")
+		viper.Reset()
+		viper.Set(c.ViperKeySFTPKeyFile, file.Name())
+
+		actual, err := LoadConfig(s.fs)
+
+		s.NoError(err)
+		s.Equal(file.Name(), actual.GetSFTPKeyFile())
+		s.Equal([]byte("private key"), actual.GetSFTPKey())
+	})
 }

pkg/constants/constants.go

@@ -1,23 +1,29 @@
 package constants
 
 const (
-	ViperKeyBindAddress = "bind_address"
-	ViperKeyPort        = "port"
-	ViperKeyLogLevel    = "log_level"
-	ViperKeySFTPHost    = "sftp_host"
-	ViperKeySFTPPort    = "sftp_port"
-	ViperKeySFTPUser    = "sftp_user"
-	ViperKeySFTPPass    = "sftp_pass"
+	ViperKeyBindAddress       = "bind_address"
+	ViperKeyPort              = "port"
+	ViperKeyLogLevel          = "log_level"
+	ViperKeySFTPHost          = "sftp_host"
+	ViperKeySFTPPort          = "sftp_port"
+	ViperKeySFTPUser          = "sftp_user"
+	ViperKeySFTPPass          = "sftp_pass"
+	ViperKeySFTPKey           = "sftp_key"
+	ViperKeySFTPKeyFile       = "sftp_key_file"
+	ViperKeySFTPKeyPassphrase = "sftp_key_passphrase"
 )
 
 const (
-	FlagBindAddress = "bind-address"
-	FlagPort        = "port"
-	FlagLogLevel    = "log-level"
-	FlagSFTPHost    = "sftp-host"
-	FlagSFTPPort    = "sftp-port"
-	FlagSFTPUser    = "sftp-user"
-	FlagSFTPPass    = "sftp-pass"
+	FlagBindAddress       = "bind-address"
+	FlagPort              = "port"
+	FlagLogLevel          = "log-level"
+	FlagSFTPHost          = "sftp-host"
+	FlagSFTPPort          = "sftp-port"
+	FlagSFTPUser          = "sftp-user"
+	FlagSFTPPass          = "sftp-pass"
+	FlagSFTPKey           = "sftp-key"
+	FlagSFTPKeyFile       = "sftp-key-file"
+	FlagSFTPKeyPassphrase = "sftp-key-passphrase"
 )
 
 const Namespace = "sftp"

pkg/utils/utils.go

@@ -0,0 +1,7 @@
+package utils
+
+import "strings"
+
+func IsNotEmpty(s string) bool {
+	return strings.TrimSpace(s) != ""
+}