/*
    Copyright (C) 2020 Accurics, 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 vulnerability

import (
	"context"
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/accurics/terrascan/pkg/iac-providers/output"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/ecr"
)

type mockECRScanner struct {
	client                          *ecr.ECR
	scanfindings                    *ecr.DescribeImageScanFindingsOutput
	scanfindingError                error
	waitUntilImageScanCompleteError error
	startImageScanOutput            *ecr.StartImageScanOutput
	startImageScanError             error
}

func (m mockECRScanner) newClient() *ecr.ECR {
	return m.client
}

func (m mockECRScanner) describeImageScanFindingsWithContext(ctx context.Context, client *ecr.ECR, input *ecr.DescribeImageScanFindingsInput, opts ...request.Option) (*ecr.DescribeImageScanFindingsOutput, error) {
	return m.scanfindings, m.scanfindingError
}

func (m mockECRScanner) waitUntilImageScanComplete(client *ecr.ECR, input *ecr.DescribeImageScanFindingsInput) error {
	return m.waitUntilImageScanCompleteError
}

func (m mockECRScanner) startImageScanWithContext(ctx context.Context, client *ecr.ECR, input *ecr.StartImageScanInput, opts ...request.Option) (*ecr.StartImageScanOutput, error) {
	return m.startImageScanOutput, m.startImageScanError
}

func TestGetImageScanResult(t *testing.T) {
	complete := "COMPLETE"
	fail := "FAILED"
	unknown := "UNKNOWN"
	time := time.Now()
	type fields struct {
		scanner ecrScanner
	}
	type args struct {
		ctx          context.Context
		client       *ecr.ECR
		image        string
		imageDetails ImageDetails
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    *ecr.DescribeImageScanFindingsOutput
		wantErr error
	}{
		{
			name: "error finding scan vulnerability",
			fields: fields{
				scanner: mockECRScanner{
					scanfindingError: fmt.Errorf("error finding vulnerability"),
				},
			},
			args: args{
				ctx:    context.Background(),
				client: &ecr.ECR{},
				image:  "24557.dkr.ecr.us-east-2.amazonaws.com/test:test",
				imageDetails: ImageDetails{
					Tag:        "test",
					Repository: "test",
					Registry:   "24557.dkr.ecr.us-east-2.amazonaws.com",
				},
			},
			wantErr: fmt.Errorf("error finding vulnerability"),
			want:    nil,
		},
		{
			name: "found scan finding with status completed",
			fields: fields{
				scanner: mockECRScanner{
					scanfindings: &ecr.DescribeImageScanFindingsOutput{
						ImageScanStatus: &ecr.ImageScanStatus{
							Status: &complete,
						},
						ImageScanFindings: &ecr.ImageScanFindings{
							ImageScanCompletedAt: &time,
						},
					},
				},
			},
			args: args{
				ctx:    context.Background(),
				client: &ecr.ECR{},
				image:  "24557.dkr.ecr.us-east-2.amazonaws.com/test:test",
				imageDetails: ImageDetails{
					Tag:        "test",
					Repository: "test",
					Registry:   "24557.dkr.ecr.us-east-2.amazonaws.com",
				},
			},
			wantErr: nil,
			want: &ecr.DescribeImageScanFindingsOutput{
				ImageScanStatus: &ecr.ImageScanStatus{
					Status: &complete,
				},
				ImageScanFindings: &ecr.ImageScanFindings{
					ImageScanCompletedAt: &time,
				},
			},
		},
		{
			name: "found scan finding with status fail",
			fields: fields{
				scanner: mockECRScanner{
					scanfindings: &ecr.DescribeImageScanFindingsOutput{
						ImageScanStatus: &ecr.ImageScanStatus{
							Status: &fail,
						},
						ImageScanFindings: &ecr.ImageScanFindings{
							ImageScanCompletedAt: &time,
						},
					},
				},
			},
			args: args{
				ctx:    context.Background(),
				client: &ecr.ECR{},
				image:  "24557.dkr.ecr.us-east-2.amazonaws.com/test:test",
				imageDetails: ImageDetails{
					Tag:        "test",
					Repository: "test",
					Registry:   "24557.dkr.ecr.us-east-2.amazonaws.com",
				},
			},
			wantErr: nil,
			want: &ecr.DescribeImageScanFindingsOutput{
				ImageScanStatus: &ecr.ImageScanStatus{
					Status: &fail,
				},
				ImageScanFindings: &ecr.ImageScanFindings{
					ImageScanCompletedAt: &time,
				},
			},
		},
		{
			name: "found scan finding with status unknown",
			fields: fields{
				scanner: mockECRScanner{
					scanfindings: &ecr.DescribeImageScanFindingsOutput{
						ImageScanStatus: &ecr.ImageScanStatus{
							Status: &unknown,
						},
						ImageScanFindings: &ecr.ImageScanFindings{
							ImageScanCompletedAt: &time,
						},
					},
				},
			},
			args: args{
				ctx:    context.Background(),
				client: &ecr.ECR{},
				image:  "24557.dkr.ecr.us-east-2.amazonaws.com/test:test",
				imageDetails: ImageDetails{
					Tag:        "test",
					Repository: "test",
					Registry:   "24557.dkr.ecr.us-east-2.amazonaws.com",
				},
			},
			wantErr: nil,
			want: &ecr.DescribeImageScanFindingsOutput{
				ImageScanStatus: &ecr.ImageScanStatus{
					Status: &unknown,
				},
				ImageScanFindings: &ecr.ImageScanFindings{
					ImageScanCompletedAt: &time,
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := &ECR{
				scanner: tt.fields.scanner,
			}
			got, err := e.GetImageScanResult(tt.args.ctx, tt.args.client, tt.args.image, tt.args.imageDetails)
			if !reflect.DeepEqual(err, tt.wantErr) {
				t.Errorf("ECR.GetImageScanResult() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("ECR.GetImageScanResult() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestStartImageScan(t *testing.T) {
	type fields struct {
		scanner ecrScanner
	}
	type args struct {
		ctx          context.Context
		client       *ecr.ECR
		image        string
		imageDetails ImageDetails
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		wantErr error
	}{
		{
			name: "error start image scan",
			fields: fields{
				scanner: mockECRScanner{
					startImageScanError: fmt.Errorf("error starting image scan"),
				},
			},
			args: args{
				ctx:    context.Background(),
				client: &ecr.ECR{},
				image:  "24557.dkr.ecr.us-east-2.amazonaws.com/test:test",
				imageDetails: ImageDetails{
					Tag:        "test",
					Repository: "test",
					Registry:   "24557.dkr.ecr.us-east-2.amazonaws.com",
				},
			},
			wantErr: fmt.Errorf("error starting image scan"),
		},
		{
			name: "error while image scan to complete",
			fields: fields{
				scanner: mockECRScanner{
					waitUntilImageScanCompleteError: fmt.Errorf("error image scan waiting"),
				},
			},
			args: args{
				ctx:    context.Background(),
				client: &ecr.ECR{},
				image:  "24557.dkr.ecr.us-east-2.amazonaws.com/test:test",
				imageDetails: ImageDetails{
					Tag:        "test",
					Repository: "test",
					Registry:   "24557.dkr.ecr.us-east-2.amazonaws.com",
				},
			},
			wantErr: fmt.Errorf("error image scan waiting"),
		},
		{
			name: "image scan started and complted successfully",
			fields: fields{
				scanner: mockECRScanner{},
			},
			args: args{
				ctx:    context.Background(),
				client: &ecr.ECR{},
				image:  "24557.dkr.ecr.us-east-2.amazonaws.com/test:test",
				imageDetails: ImageDetails{
					Tag:        "test",
					Repository: "test",
					Registry:   "24557.dkr.ecr.us-east-2.amazonaws.com",
				},
			},
			wantErr: nil,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := &ECR{
				scanner: tt.fields.scanner,
			}
			err := e.StartImageScan(tt.args.ctx, tt.args.client, tt.args.image, tt.args.imageDetails)
			if !reflect.DeepEqual(err, tt.wantErr) {
				t.Errorf("ECR.StartImageScan() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func TestECRScanImage(t *testing.T) {
	complete := "COMPLETE"
	time := time.Now()
	type fields struct {
		scanner ecrScanner
	}
	type args struct {
		ctx   context.Context
		image string
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    *ecr.DescribeImageScanFindingsOutput
		wantErr bool
	}{
		{
			name: "inavlid image reference",
			fields: fields{
				scanner: mockECRScanner{},
			},
			args: args{
				ctx:   context.Background(),
				image: "test",
			},
			want:    nil,
			wantErr: true,
		},
		{
			name: "valid image reference",
			fields: fields{
				scanner: mockECRScanner{
					scanfindings: &ecr.DescribeImageScanFindingsOutput{
						ImageScanStatus: &ecr.ImageScanStatus{
							Status: &complete,
						},
						ImageScanFindings: &ecr.ImageScanFindings{
							ImageScanCompletedAt: &time,
						},
					},
				},
			},
			args: args{
				ctx:   context.Background(),
				image: "24557.dkr.ecr.us-east-2.amazonaws.com/test",
			},
			want: &ecr.DescribeImageScanFindingsOutput{
				ImageScanStatus: &ecr.ImageScanStatus{
					Status: &complete,
				},
				ImageScanFindings: &ecr.ImageScanFindings{
					ImageScanCompletedAt: &time,
				},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := &ECR{
				scanner: tt.fields.scanner,
			}
			got, err := e.ScanImage(tt.args.ctx, tt.args.image)
			if (err != nil) != tt.wantErr {
				t.Errorf("ECR.ScanImage() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("ECR.ScanImage() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestECRGetVulnerabilities(t *testing.T) {
	complete := "COMPLETE"
	severityLow := "Low"
	time := time.Now()
	type fields struct {
		scanner ecrScanner
	}
	type args struct {
		container output.ContainerDetails
		options   map[string]interface{}
	}
	tests := []struct {
		name                string
		fields              fields
		args                args
		wantVulnerabilities []output.Vulnerability
	}{
		{
			name: "invalid image reference",
			fields: fields{
				scanner: mockECRScanner{},
			},
			args: args{
				container: output.ContainerDetails{
					Image: "test",
				},
			},
			wantVulnerabilities: nil,
		},
		{
			name: "found image vulnerabilities",
			fields: fields{
				scanner: mockECRScanner{
					scanfindings: &ecr.DescribeImageScanFindingsOutput{
						ImageScanStatus: &ecr.ImageScanStatus{
							Status: &complete,
						},
						ImageScanFindings: &ecr.ImageScanFindings{
							ImageScanCompletedAt: &time,
							Findings: []*ecr.ImageScanFinding{
								{
									Severity: &severityLow,
								},
							},
						},
					},
				},
			},
			args: args{
				container: output.ContainerDetails{
					Image: "24557.dkr.ecr.us-east-2.amazonaws.com/test",
				},
			},
			wantVulnerabilities: []output.Vulnerability{
				{
					PrimaryURL: "https://cve.mitre.org/cgi-bin/cvename.cgi?name=",
					Severity:   "LOW",
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := &ECR{
				scanner: tt.fields.scanner,
			}
			if gotVulnerabilities := e.getVulnerabilities(tt.args.container, tt.args.options); !reflect.DeepEqual(gotVulnerabilities, tt.wantVulnerabilities) {
				t.Errorf("ECR.getVulnerabilities() = %v, want %v", gotVulnerabilities, tt.wantVulnerabilities)
			}
		})
	}
}
