/*
    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"

	"github.com/Azure/azure-sdk-for-go/profiles/preview/preview/containerregistry/runtime/containerregistry"
	"github.com/Azure/azure-sdk-for-go/services/resourcegraph/mgmt/2021-03-01/resourcegraph"
	"github.com/Azure/go-autorest/autorest"
	"github.com/accurics/terrascan/pkg/iac-providers/output"
)

type mockScanner struct {
	basicAuthorizer *autorest.BasicAuthorizer
	queryError      error
	authorizerError error
	attributesError error
	authorizer      autorest.Authorizer
	tagClient       containerregistry.TagClient
	queryResponse   resourcegraph.QueryResponse
	attributes      containerregistry.TagAttributes
	validImage      bool
	digest          string
}

//newTagClient
func (m mockScanner) newTagClient(loginURL string) containerregistry.TagClient {
	return m.tagClient
}

//newBasicAuthorizer
func (m mockScanner) newBasicAuthorizer(username string, password string) *autorest.BasicAuthorizer {
	return m.basicAuthorizer
}

//newAuthorizerFromFile
func (m mockScanner) newAuthorizerFromFile() (autorest.Authorizer, error) {
	return m.authorizer, m.authorizerError
}

//newResourcegraphClient
func (mockScanner) newResourcegraphClient() resourcegraph.BaseClient {
	return resourcegraph.New()
}

//getResources
func (m mockScanner) getResources(ctx context.Context, request resourcegraph.QueryRequest, client resourcegraph.BaseClient) (resourcegraph.QueryResponse, error) {
	return m.queryResponse, m.queryError
}

//GetAttributes
func (m mockScanner) getAttributes(ctx context.Context, name string, reference string, client containerregistry.TagClient) (result containerregistry.TagAttributes, err error) {
	return m.attributes, m.attributesError
}

func TestScanImage(t *testing.T) {
	type fields struct {
		scanner acrScanner
	}
	type args struct {
		image string
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		want    []output.ACRResponse
		wantErr error
	}{
		{
			name: "error getting authorizer",
			fields: fields{
				scanner: mockScanner{
					authorizer:      nil,
					authorizerError: fmt.Errorf("error creating authorizer"),
				},
			},
			args: args{
				image: "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57",
			},
			want:    []output.ACRResponse{},
			wantErr: fmt.Errorf("error creating authorizer"),
		},
		{
			name: "invalid image",
			fields: fields{
				scanner: mockScanner{
					authorizer: nil,
					validImage: false,
				},
			},
			args: args{
				image: "test",
			},
			want:    []output.ACRResponse{},
			wantErr: fmt.Errorf("invalid image reference test "),
		},
		{
			name: "error from query execution",
			fields: fields{
				scanner: mockScanner{
					queryError: fmt.Errorf("error finding resource"),
					validImage: true,
				},
			},
			args: args{
				image: "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57",
			},
			want:    []output.ACRResponse{},
			wantErr: fmt.Errorf("error finding resource"),
		},
		{
			name: "unmarshal error",
			fields: fields{
				scanner: mockScanner{
					queryError: nil,
					validImage: true,
					queryResponse: resourcegraph.QueryResponse{
						Data: "nil",
					},
				},
			},
			args: args{
				image: "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57",
			},
			want:    []output.ACRResponse{},
			wantErr: fmt.Errorf("error unmarshaling image %s scan results", "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57"),
		},
		{
			name: "marshal error",
			fields: fields{
				scanner: mockScanner{
					queryError: nil,
					validImage: true,
					queryResponse: resourcegraph.QueryResponse{
						Data: make(chan int),
					},
				},
			},
			args: args{
				image: "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57",
			},
			want:    []output.ACRResponse{},
			wantErr: fmt.Errorf("error marshaling image %s scan results", "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57"),
		},
		{
			name: "valid response from query execution",
			fields: fields{
				scanner: mockScanner{
					queryError: nil,
					validImage: true,
					queryResponse: resourcegraph.QueryResponse{
						Data: []output.ACRResponse{{
							Properties: output.ACRVulnerabilityConfig{
								ID: "test",
							},
						}},
					},
				},
			},
			args: args{
				image: "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57",
			},
			want: []output.ACRResponse{
				{
					Properties: output.ACRVulnerabilityConfig{
						ID: "test",
					},
				},
			},
			wantErr: nil,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			a := &ACR{
				scanner: tt.fields.scanner,
			}
			got, err := a.ScanImage(tt.args.image)
			if !reflect.DeepEqual(err, tt.wantErr) {
				t.Errorf("ACR.ScanImage() got error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("ACR.ScanImage() got = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestGetACRImageDetails(t *testing.T) {
	digest := "sha256:3f13b4376446cf92b0cb9a"
	type fields struct {
		scanner acrScanner
	}
	type args struct {
		image     string
		authrizer autorest.Authorizer
	}
	tests := []struct {
		name             string
		fields           fields
		args             args
		wantImageDetails ImageDetails
		wantValidImage   bool
	}{
		{
			name: "invalid image",
			fields: fields{
				scanner: nil,
			},
			args: args{
				image: "test",
			},
			wantImageDetails: ImageDetails{
				Registry: "test",
			},
			wantValidImage: false,
		},
		{
			name: "unable to find image digest",
			fields: fields{
				scanner: mockScanner{
					digest: "",
				},
			},
			args: args{
				image: "test.azurecr.io/nginx",
			},
			wantImageDetails: ImageDetails{
				Registry:   "test.azurecr.io",
				Repository: "nginx",
				Tag:        "latest",
			},
			wantValidImage: false,
		},
		{
			name: "found image digest",
			fields: fields{
				scanner: mockScanner{
					digest: "sha256:3f13b4376446cf92b0cb9a",
					attributes: containerregistry.TagAttributes{
						Attributes: &containerregistry.TagAttributesBase{
							Digest: &digest,
						},
					},
				},
			},
			args: args{
				image: "test.azurecr.io/nginx",
			},
			wantImageDetails: ImageDetails{
				Registry:   "test.azurecr.io",
				Repository: "nginx",
				Tag:        "latest",
				Digest:     "sha256:3f13b4376446cf92b0cb9a",
			},
			wantValidImage: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			a := &ACR{
				scanner: tt.fields.scanner,
			}
			gotImageDetails, gotValidImage := a.getACRImageDetails(tt.args.image, tt.args.authrizer)
			if !reflect.DeepEqual(gotImageDetails, tt.wantImageDetails) {
				t.Errorf("ACR.getACRImageDetails() gotImageDetails = %v, want %v", gotImageDetails, tt.wantImageDetails)
			}
			if gotValidImage != tt.wantValidImage {
				t.Errorf("ACR.getACRImageDetails() gotValidImage = %v, want %v", gotValidImage, tt.wantValidImage)
			}
		})
	}
}

func TestGetImageDigestFromTag(t *testing.T) {
	digest := "sha256:3f13b4376446cf92b0cb9a"
	type fields struct {
		scanner acrScanner
	}
	type args struct {
		image        string
		imageDetails ImageDetails
		loginURL     string
	}
	tests := []struct {
		name       string
		fields     fields
		args       args
		wantDigest string
	}{
		{
			name: "error getting digest",
			fields: fields{
				scanner: mockScanner{
					attributesError: fmt.Errorf("error getting image attributes"),
				},
			},
			args: args{
				image: "test.azurecr.io/nginx:latest",
				imageDetails: ImageDetails{
					Registry:   "test.azurecr.io",
					Repository: "nginx",
					Tag:        "latest",
				},
				loginURL: "https://test.azurecr.io",
			},
			wantDigest: "",
		},
		{
			name: "digest found",
			fields: fields{
				scanner: mockScanner{
					attributes: containerregistry.TagAttributes{
						Attributes: &containerregistry.TagAttributesBase{
							Digest: &digest,
						},
					},
				},
			},
			args: args{
				image: "test.azurecr.io/nginx:latest",
				imageDetails: ImageDetails{
					Registry:   "test.azurecr.io",
					Repository: "nginx",
					Tag:        "latest",
				},
				loginURL: "https://test.azurecr.io",
			},
			wantDigest: digest,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			a := &ACR{
				scanner: tt.fields.scanner,
			}
			if gotDigest := a.getImageDigestFromTag(tt.args.image, tt.args.imageDetails, tt.args.loginURL); gotDigest != tt.wantDigest {
				t.Errorf("ACR.getImageDigestFromTag() = %v, want %v", gotDigest, tt.wantDigest)
			}
		})
	}
}

func TestGetVulnerabilities(t *testing.T) {
	type fields struct {
		scanner acrScanner
	}
	type args struct {
		container output.ContainerDetails
		options   map[string]interface{}
	}
	tests := []struct {
		name                string
		fields              fields
		args                args
		wantVulnerabilities []output.Vulnerability
	}{
		{
			name: "error scanning image",
			args: args{
				container: output.ContainerDetails{
					Image: "test",
				},
			},
			fields: fields{
				scanner: mockScanner{
					authorizer:      nil,
					authorizerError: fmt.Errorf("error creating authorizer"),
				},
			},
			wantVulnerabilities: nil,
		},
		{
			name: "found vulnerability in image",
			args: args{
				container: output.ContainerDetails{
					Image: "test.azurecr.io/nginx:latest@sha256:3f13b4376446cf92b0cb9a5c46ba75d57",
				},
			},
			fields: fields{
				scanner: mockScanner{
					queryError: nil,
					validImage: true,
					queryResponse: resourcegraph.QueryResponse{
						Data: []output.ACRResponse{{
							Properties: output.ACRVulnerabilityConfig{
								ID: "test",
								AdditionalData: output.ACRAdditionalData{
									Cve: []output.ACRCve{
										{
											Link:  "http://link",
											Title: "CVE-6767",
										},
									},
								},
							},
						}},
					},
				},
			},
			wantVulnerabilities: []output.Vulnerability{
				{
					VulnerabilityID: "CVE-6767",
					PrimaryURL:      "http://link",
					CVSS: output.VendorCVSS{
						"nvd": output.CVSS{},
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			a := &ACR{
				scanner: tt.fields.scanner,
			}
			if gotVulnerabilities := a.getVulnerabilities(tt.args.container, tt.args.options); !reflect.DeepEqual(gotVulnerabilities, tt.wantVulnerabilities) {
				t.Errorf("ACR.getVulnerabilities() = %#v, want %v", gotVulnerabilities, tt.wantVulnerabilities)
			}
		})
	}
}
