/*
    Copyright (C) 2022 Tenable, Inc.

	Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

		http://www.apache.org/licenses/LICENSE-2.0

	Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
*/

package output

import (
	"reflect"
	"testing"

	"github.com/aws/aws-sdk-go/service/ecr"
	grafeaspb "google.golang.org/genproto/googleapis/grafeas/v1"
)

func TestCVSSPrepareFromECRImageScanAttribute(t *testing.T) {
	v2VectorKey := ecrCvss + "2_" + ecrCvssVector
	v2ScoreKey := ecrCvss + "2_" + ecrCvssScore
	v3VectorKey := ecrCvss + "3_" + ecrCvssVector
	v3ScoreKey := ecrCvss + "3_" + ecrCvssScore
	value := "6.7"

	type args struct {
		attribute *ecr.Attribute
	}
	tests := []struct {
		name string
		args args
		want *CVSS
	}{
		{
			name: "v2 vector present in cvss object",
			args: args{
				attribute: &ecr.Attribute{
					Key:   &v2VectorKey,
					Value: &value,
				},
			},
			want: &CVSS{
				V2Vector: "6.7",
			},
		},
		{
			name: "v2 score present in cvss object",
			args: args{
				attribute: &ecr.Attribute{
					Key:   &v2ScoreKey,
					Value: &value,
				},
			},
			want: &CVSS{
				V2Score: 6.7,
			},
		},
		{
			name: "v3 score present in cvss object",
			args: args{
				attribute: &ecr.Attribute{
					Key:   &v3ScoreKey,
					Value: &value,
				},
			},
			want: &CVSS{
				V3Score: 6.7,
			},
		},
		{
			name: "v3 vector present in cvss object",
			args: args{
				attribute: &ecr.Attribute{
					Key:   &v3VectorKey,
					Value: &value,
				},
			},
			want: &CVSS{
				V3Vector: "6.7",
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cvss := &CVSS{}
			cvss.PrepareFromECRImageScanAttribute(tt.args.attribute)
			if !reflect.DeepEqual(cvss, tt.want) {
				t.Errorf("PrepareFromECRImageScanAttribute() got = %v, want %v", cvss, tt.want)
				return
			}
		})
	}
}

func TestVulnerabilityPrepareFromECRImageScan(t *testing.T) {

	name := "test_vulnerability"
	severity := "low"
	description := "description"
	uri := ""
	packageName := "test_package"
	packageNameKey := ecrPackageName
	packageVersion := "v1"
	packageVersionKey := ecrPackageVersion
	v3ScoreKey := ecrCvss + "3_" + ecrCvssScore
	value := "6.7"

	type args struct {
		imageScanFinding *ecr.ImageScanFinding
	}
	tests := []struct {
		name string
		args args
		want *Vulnerability
	}{
		{
			name: "get Vulnerability object ready",
			args: args{
				imageScanFinding: &ecr.ImageScanFinding{
					Description: &description,
					Name:        &name,
					Severity:    &severity,
					Uri:         &uri,
					Attributes: []*ecr.Attribute{
						{
							Key:   &packageNameKey,
							Value: &packageName,
						},
						{
							Key:   &packageVersionKey,
							Value: &packageVersion,
						},
						{
							Key:   &v3ScoreKey,
							Value: &value,
						},
						{
							Key:   nil,
							Value: nil,
						},
					},
				},
			},
			want: &Vulnerability{
				VulnerabilityID:  "test_vulnerability",
				PkgName:          "test_package",
				InstalledVersion: "v1",
				PrimaryURL:       "https://cve.mitre.org/cgi-bin/cvename.cgi?name=test_vulnerability",
				Description:      "description",
				Severity:         "LOW",
				CVSS: VendorCVSS{
					"nvd": CVSS{
						V3Score: 6.7,
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			v := &Vulnerability{}
			v.PrepareFromECRImageScan(tt.args.imageScanFinding)
			if !reflect.DeepEqual(v, tt.want) {
				t.Errorf("PrepareFromECRImageScan() got = %v, want %v", v, tt.want)
				return
			}
		})
	}
}

func TestVulnerabilityPrepareFromACRImageScan(t *testing.T) {
	type args struct {
		acrResponse ACRResponse
	}
	tests := []struct {
		name string
		args args
		want *Vulnerability
	}{
		{
			name: "vulnerability not found",
			args: args{
				acrResponse: ACRResponse{},
			},
			want: &Vulnerability{},
		},
		{
			name: "vulnerability object from acr response",
			args: args{
				acrResponse: ACRResponse{
					Properties: ACRVulnerabilityConfig{
						ID:          "test_ID",
						Description: "description",
						AdditionalData: ACRAdditionalData{
							Cve: []ACRCve{
								{
									Title: "test_Title",
									Link:  "",
								},
							},
							Cvss: ACRCvss{
								V2: ACRCvssBase{
									Base: 4.5,
								},
								V3: ACRCvssBase{
									Base: 4.5,
								},
							},
						},
						Status: ACRStatus{
							Severity: "low",
						},
					},
				},
			},
			want: &Vulnerability{
				VulnerabilityID: "test_Title",
				PrimaryURL:      "https://cve.mitre.org/cgi-bin/cvename.cgi?name=test_Title",
				Description:     "description",
				Severity:        "LOW",
				CVSS: VendorCVSS{
					"nvd": CVSS{V2Score: 4.5, V3Score: 4.5}}},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			v := &Vulnerability{}
			v.PrepareFromACRImageScan(tt.args.acrResponse)
			if !reflect.DeepEqual(v, tt.want) {
				t.Errorf("PrepareFromACRImageScan() got = %v, want %v", v, tt.want)
				return
			}
		})
	}
}

func TestVulnerabilityPrepareFromGCRImageScan(t *testing.T) {

	type args struct {
		gcpVulnerability *grafeaspb.Occurrence
	}
	tests := []struct {
		name string
		args args
		want *Vulnerability
	}{
		{
			name: "complete vulnerability object",
			args: args{
				gcpVulnerability: &grafeaspb.Occurrence{
					NoteName: "https://test/CVE-1234",
					Details: &grafeaspb.Occurrence_Vulnerability{
						Vulnerability: &grafeaspb.VulnerabilityOccurrence{
							Severity:  2,
							CvssScore: 3.3,
							Type:      "os",
							PackageIssue: []*grafeaspb.VulnerabilityOccurrence_PackageIssue{
								{
									AffectedPackage: "test",
									AffectedVersion: &grafeaspb.Version{
										FullName: "v1.4.3",
									},
									FixedVersion: &grafeaspb.Version{
										FullName: "v1.4.3",
									},
								},
							},
							LongDescription:  "long description",
							ShortDescription: "short description",
							RelatedUrls: []*grafeaspb.RelatedUrl{
								{
									Url: "http://test_url",
								},
							},
						},
					},
				},
			},
			want: &Vulnerability{
				Type:             "os",
				VulnerabilityID:  "CVE-1234",
				PkgName:          "test",
				InstalledVersion: "v1.4.3",
				FixedVersion:     "v1.4.3",
				PrimaryURL:       "https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-1234",
				Description:      "long description",
				Severity:         "LOW",
				CVSS: VendorCVSS{
					"nvd": CVSS{V2Score: 3.299999952316284}},
				References: []string{"http://test_url"},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			v := &Vulnerability{}
			v.PrepareFromGCRImageScan(tt.args.gcpVulnerability)
			if !reflect.DeepEqual(v, tt.want) {
				t.Errorf("PrepareFromGCRImageScan() got = %v, want %v", v, tt.want)
				return
			}
		})
	}
}

func TestVulnerability_PrepareFromHarborImageScan(t *testing.T) {

	type args struct {
		vulnerability map[string]interface{}
	}
	tests := []struct {
		name string
		args args
		want *Vulnerability
	}{
		{
			name: "complete vulnerability object",
			args: args{
				vulnerability: map[string]interface{}{
					"id":          "CVE-2021-36159",
					"package":     "apk-tools",
					"version":     "2.10.5-r1",
					"fix_version": "2.10.7-r0",
					"severity":    "Critical",
					"description": "libfetch before 2021-07-26, as used",
					"vendor_attributes": map[string]interface{}{
						"CVSS": map[string]interface{}{
							"nvd": map[string]interface{}{
								"V2Score":  6.4,
								"V2Vector": "AV:N/AC:L/Au:N/C:P/I:N/A:P",
								"V3Score":  9.1,
								"V3Vector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:N/A:H",
							},
						},
					},
				},
			},
			want: &Vulnerability{
				VulnerabilityID:  "CVE-2021-36159",
				PkgName:          "apk-tools",
				InstalledVersion: "2.10.5-r1",
				FixedVersion:     "2.10.7-r0",
				PrimaryURL:       "https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-36159",
				Description:      "libfetch before 2021-07-26, as used",
				Severity:         "Critical",
				CVSS: VendorCVSS{
					"nvd": CVSS{
						V2Vector: "AV:N/AC:L/Au:N/C:P/I:N/A:P",
						V3Vector: "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:N/A:H",
						V2Score:  6.4,
						V3Score:  9.1,
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			v := &Vulnerability{}
			v.PrepareFromHarborImageScan(tt.args.vulnerability)
			if !reflect.DeepEqual(v, tt.want) {
				t.Errorf("PrepareFromHarborImageScan() got = %v, want %v", v, tt.want)
				return
			}
		})
	}
}
