From 5076326589fe7a2b4a7d87b36822912b22ae064e Mon Sep 17 00:00:00 2001
From: Kota Kanbe <kotakanbe@gmail.com>
Date: Tue, 10 Apr 2018 11:53:11 +0900
Subject: [PATCH] Fix Amazon Linux 2 scanning (#630)

* fix(amazon2): fix OS version parse error
---
 config/config.go      |  8 ++++++++
 config/config_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/config/config.go b/config/config.go
index 82af5922..addfd920 100644
--- a/config/config.go
+++ b/config/config.go
@@ -664,6 +664,14 @@ func (l Distro) String() string {
 
 // MajorVersion returns Major version
 func (l Distro) MajorVersion() (ver int, err error) {
+	if l.Family == Amazon {
+		ss := strings.Fields(l.Release)
+		if len(ss) == 1 {
+			return 1, nil
+		}
+		ver, err = strconv.Atoi(ss[0])
+		return
+	}
 	if 0 < len(l.Release) {
 		ver, err = strconv.Atoi(strings.Split(l.Release, ".")[0])
 	} else {
diff --git a/config/config_test.go b/config/config_test.go
index ed7d9220..b7d61191 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -62,3 +62,42 @@ func TestSyslogConfValidate(t *testing.T) {
 		}
 	}
 }
+
+func TestMajorVersion(t *testing.T) {
+	var tests = []struct {
+		in  Distro
+		out int
+	}{
+		{
+			in: Distro{
+				Family:  Amazon,
+				Release: "2 (2017.12)",
+			},
+			out: 2,
+		},
+		{
+			in: Distro{
+				Family:  Amazon,
+				Release: "2017.12",
+			},
+			out: 1,
+		},
+		{
+			in: Distro{
+				Family:  CentOS,
+				Release: "7.10",
+			},
+			out: 7,
+		},
+	}
+
+	for i, tt := range tests {
+		ver, err := tt.in.MajorVersion()
+		if err != nil {
+			t.Errorf("[%d] err occurred: %s", i, err)
+		}
+		if tt.out != ver {
+			t.Errorf("[%d] expected %d, actual %d", i, tt.out, ver)
+		}
+	}
+}