Merge pull request #686 from nicjohnson145/feat--device-allowlist

feat: create allow-list for filtering down devices to only a subset
AnalogJ · Sep 20, 2024 · e0fe17a · e0fe17a
2 parents 48feee5 + c9429c6
commit e0fe17a
Show file tree

Hide file tree

Showing 7 changed files with 116 additions and 1 deletion.
diff --git a/collector/pkg/config/config.go b/collector/pkg/config/config.go
@@ -20,7 +20,7 @@ import (
 type configuration struct {
 	*viper.Viper
 
-	deviceOverrides []models.ScanOverride
+	deviceOverrides    []models.ScanOverride
 }
 
 //Viper uses the following precedence order. Each item takes precedence over the item below it:
@@ -50,6 +50,8 @@ func (c *configuration) Init() error {
 
 	//c.SetDefault("collect.short.command", "-a -o on -S on")
 
+	c.SetDefault("allow_listed_devices", []string{})
+
 	//if you want to load a non-standard location system config file (~/drawbridge.yml), use ReadConfig
 	c.SetConfigType("yaml")
 	//c.SetConfigName("drawbridge")
@@ -186,3 +188,18 @@ func (c *configuration) GetCommandMetricsSmartArgs(deviceName string) string {
 	}
 	return c.GetString("commands.metrics_smart_args")
 }
+
+func (c *configuration) IsAllowlistedDevice(deviceName string) bool {
+	allowList := c.GetStringSlice("allow_listed_devices")
+	if len(allowList) == 0 {
+		return true
+	}
+
+	for _, item := range allowList {
+		if item == deviceName {
+			return true
+		}
+	}
+
+	return false
+}
diff --git a/collector/pkg/config/config_test.go b/collector/pkg/config/config_test.go
@@ -144,3 +144,29 @@ func TestConfiguration_OverrideDeviceCommands_MetricsInfoArgs(t *testing.T) {
 	require.Equal(t, "--info --json", testConfig.GetCommandMetricsInfoArgs("/dev/sdb"))
 	//require.Equal(t, []models.ScanOverride{{Device: "/dev/sda", DeviceType: nil, Commands: {MetricsInfoArgs: "--info --json -T "}}}, scanOverrides)
 }
+
+func TestConfiguration_DeviceAllowList(t *testing.T) {
+	t.Parallel()
+
+	t.Run("present", func(t *testing.T) {
+		testConfig, err := config.Create()
+		require.NoError(t, err)
+
+		require.NoError(t, testConfig.ReadConfig(path.Join("testdata", "allow_listed_devices_present.yaml")))
+
+		require.True(t, testConfig.IsAllowlistedDevice("/dev/sda"), "/dev/sda should be allow listed")
+		require.False(t, testConfig.IsAllowlistedDevice("/dev/sdc"), "/dev/sda should not be allow listed")
+	})
+
+	t.Run("missing", func(t *testing.T) {
+		testConfig, err := config.Create()
+		require.NoError(t, err)
+
+		// Really just any other config where the key is full missing
+		require.NoError(t, testConfig.ReadConfig(path.Join("testdata", "override_device_commands.yaml")))
+
+		// Anything should be allow listed if the key isnt there
+		require.True(t, testConfig.IsAllowlistedDevice("/dev/sda"), "/dev/sda should be allow listed")
+		require.True(t, testConfig.IsAllowlistedDevice("/dev/sdc"), "/dev/sda should be allow listed")
+	})
+}
diff --git a/collector/pkg/config/interface.go b/collector/pkg/config/interface.go
@@ -25,4 +25,6 @@ type Interface interface {
 	GetDeviceOverrides() []models.ScanOverride
 	GetCommandMetricsInfoArgs(deviceName string) string
 	GetCommandMetricsSmartArgs(deviceName string) string
+
+	IsAllowlistedDevice(deviceName string) bool
 }
diff --git a/collector/pkg/config/mock/mock_config.go b/collector/pkg/config/mock/mock_config.go
diff --git a/collector/pkg/config/testdata/allow_listed_devices_present.yaml b/collector/pkg/config/testdata/allow_listed_devices_present.yaml
@@ -0,0 +1,3 @@
+allow_listed_devices:
+- /dev/sda
+- /dev/sdb
diff --git a/collector/pkg/detect/detect.go b/collector/pkg/detect/detect.go
@@ -124,6 +124,11 @@ func (d *Detect) TransformDetectedDevices(detectedDeviceConns models.Scan) []mod
 
 		deviceFile := strings.ToLower(scannedDevice.Name)
 
+		// If the user has defined a device allow list, and this device isnt there, then ignore it
+		if !d.Config.IsAllowlistedDevice(deviceFile) {
+			continue
+		}
+
 		detectedDevice := models.Device{
 			HostId:     d.Config.GetString("host.id"),
 			DeviceType: scannedDevice.Type,

diff --git a/collector/pkg/detect/detect_test.go b/collector/pkg/detect/detect_test.go
@@ -24,6 +24,7 @@ func TestDetect_SmartctlScan(t *testing.T) {
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{})
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 
 	fakeShell := mock_shell.NewMockInterface(mockCtrl)
 	testScanResults, err := os.ReadFile("testdata/smartctl_scan_simple.json")
@@ -53,6 +54,7 @@ func TestDetect_SmartctlScan_Megaraid(t *testing.T) {
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{})
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 
 	fakeShell := mock_shell.NewMockInterface(mockCtrl)
 	testScanResults, err := os.ReadFile("testdata/smartctl_scan_megaraid.json")
@@ -85,6 +87,7 @@ func TestDetect_SmartctlScan_Nvme(t *testing.T) {
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{})
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 
 	fakeShell := mock_shell.NewMockInterface(mockCtrl)
 	testScanResults, err := os.ReadFile("testdata/smartctl_scan_nvme.json")
@@ -116,6 +119,7 @@ func TestDetect_TransformDetectedDevices_Empty(t *testing.T) {
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{})
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 
 	detectedDevices := models.Scan{
 		Devices: []models.ScanDevice{
@@ -149,6 +153,7 @@ func TestDetect_TransformDetectedDevices_Ignore(t *testing.T) {
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{{Device: "/dev/sda", DeviceType: nil, Ignore: true}})
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 
 	detectedDevices := models.Scan{
 		Devices: []models.ScanDevice{
@@ -180,6 +185,7 @@ func TestDetect_TransformDetectedDevices_Raid(t *testing.T) {
 	fakeConfig.EXPECT().GetString("host.id").AnyTimes().Return("")
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{
 		{
 			Device:     "/dev/bus/0",
@@ -223,6 +229,7 @@ func TestDetect_TransformDetectedDevices_Simple(t *testing.T) {
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{{Device: "/dev/sda", DeviceType: []string{"sat+megaraid"}}})
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 	detectedDevices := models.Scan{
 		Devices: []models.ScanDevice{
 			{
@@ -256,6 +263,7 @@ func TestDetect_TransformDetectedDevices_WithoutDeviceTypeOverride(t *testing.T)
 	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
 	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
 	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{{Device: "/dev/sda"}})
+	fakeConfig.EXPECT().IsAllowlistedDevice(gomock.Any()).AnyTimes().Return(true)
 	detectedDevices := models.Scan{
 		Devices: []models.ScanDevice{
 			{
@@ -302,6 +310,46 @@ func TestDetect_TransformDetectedDevices_WhenDeviceNotDetected(t *testing.T) {
 	require.Equal(t, "ata", transformedDevices[0].DeviceType)
 }
 
+func TestDetect_TransformDetectedDevices_AllowListFilters(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	fakeConfig := mock_config.NewMockInterface(mockCtrl)
+	fakeConfig.EXPECT().GetString("host.id").AnyTimes().Return("")
+	fakeConfig.EXPECT().GetString("commands.metrics_smartctl_bin").AnyTimes().Return("smartctl")
+	fakeConfig.EXPECT().GetString("commands.metrics_scan_args").AnyTimes().Return("--scan --json")
+	fakeConfig.EXPECT().GetDeviceOverrides().AnyTimes().Return([]models.ScanOverride{{Device: "/dev/sda", DeviceType: []string{"sat+megaraid"}}})
+	fakeConfig.EXPECT().IsAllowlistedDevice("/dev/sda").Return(true)
+	fakeConfig.EXPECT().IsAllowlistedDevice("/dev/sdb").Return(false)
+	detectedDevices := models.Scan{
+		Devices: []models.ScanDevice{
+			{
+				Name:     "/dev/sda",
+				InfoName: "/dev/sda",
+				Protocol: "ata",
+				Type:     "ata",
+			},
+			{
+				Name:     "/dev/sdb",
+				InfoName: "/dev/sdb",
+				Protocol: "ata",
+				Type:     "ata",
+			},
+		},
+	}
+
+	d := detect.Detect{
+		Config: fakeConfig,
+	}
+
+	// test
+	transformedDevices := d.TransformDetectedDevices(detectedDevices)
+
+	// assert
+	require.Equal(t, 1, len(transformedDevices))
+	require.Equal(t, "sda", transformedDevices[0].DeviceName)
+}
+
 func TestDetect_SmartCtlInfo(t *testing.T) {
 	t.Run("should report nvme info", func(t *testing.T) {
 		ctrl := gomock.NewController(t)