package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/testing/mock_modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/pkg/agentcfg"
	"go.uber.org/zap/zaptest"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes/fake"
	clienttesting "k8s.io/client-go/testing"
)

const (
	namespace1                = "namespace1"
	namespace2                = "namespace2"
	gitlabAgentNamespace      = "gitlab-agent"
	gitlabAgentServiceAccount = "gitlab-agent"
	podName                   = "Pod name"
)

func TestScan(t *testing.T) {
	ctrl := gomock.NewController(t)
	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
	defer cancel()

	kubeClientset := fake.NewSimpleClientset()
	logParser := &CustomLogParser{}
	api := mock_modagent.NewMockApi(ctrl)
	s := scanJob{
		log:                       zaptest.NewLogger(t),
		api:                       api,
		kubeClientset:             kubeClientset,
		gitlabAgentNamespace:      gitlabAgentNamespace,
		gitlabAgentServiceAccount: gitlabAgentServiceAccount,
		agentID:                   5000000000,
		targetNamespaces:          []string{namespace1, namespace2},
		logParser:                 logParser,
		resourceRequirements:      &defaultResourceRequirements,
	}

	// Each namespace scan will return 1 vulnerability based on the response from ParsePodLogsToReport. As there are 2 namespaces this will result in 2 vulnerabilities to be transmitted to Gitlab.
	api.EXPECT().MakeGitLabRequest(gomock.Any(), "/",
		gomock.Any(),
		gomock.Any(),
	).Times(2).
		DoAndReturn(mockVulnerabilityCreate)

	api.EXPECT().MakeGitLabRequest(gomock.Any(), "/scan_result",
		gomock.Any(),
		gomock.Any(),
	).Times(1).
		DoAndReturn(func(ctx context.Context, path string, opts ...modagent.GitLabRequestOption) (*modagent.GitLabResponse, error) {
			uuids, err := getUuidsFromOptions(opts)
			require.NoError(t, err)

			// There are 2 vulnerabilities detected in total.
			// As we want to send single request for all workloads,
			// we need to check if this single request contains all returned vulnerabilities.
			expectedVulnerabiltiyIds := []string{
				"2b19bdf999f6fdd149344caf0b0ff2c9fb96fa2c7eb1eb25907b48c32980397a",
				"2b19bdf999f6fdd149344caf0b0ff2c9fb96fa2c7eb1eb25907b48c32980397a",
			}

			assert.ElementsMatch(t, expectedVulnerabiltiyIds, uuids)

			return &modagent.GitLabResponse{
				StatusCode: 200,
			}, nil
		})

	// s.Run starts a Watcher for each Trivy Pod that it creates and waits for the Pod status to change to Succeeded before processing it's logs.
	// As the Watcher is created within the scanJob.startPodScanForNamespace, I opted to use a fake watch interface to stub the Pod status change events.
	wg := startCustomWatcher(t, ctx, kubeClientset, 2, corev1.PodSucceeded)

	s.Run(ctx)
	podList, err := kubeClientset.CoreV1().Pods(gitlabAgentNamespace).List(ctx, metav1.ListOptions{})
	assert.NoError(t, err)
	// Trivy Pods should be deleted after the scan is complete
	assert.Equal(t, 0, len(podList.Items))

	wg.Wait()
}

func TestScanPodFailed(t *testing.T) {
	ctrl := gomock.NewController(t)
	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
	defer cancel()

	kubeClientset := fake.NewSimpleClientset()

	logParser := &CustomLogParser{}
	api := mock_modagent.NewMockApi(ctrl)
	s := scanJob{
		log:                       zaptest.NewLogger(t),
		api:                       api,
		kubeClientset:             kubeClientset,
		gitlabAgentNamespace:      gitlabAgentNamespace,
		gitlabAgentServiceAccount: gitlabAgentServiceAccount,
		agentID:                   5000000000,
		targetNamespaces:          []string{namespace1},
		logParser:                 logParser,
		resourceRequirements:      &defaultResourceRequirements,
	}

	api.EXPECT().MakeGitLabRequest(gomock.Any(), "/",
		gomock.Any(),
		gomock.Any(),
	).Times(0)

	api.EXPECT().MakeGitLabRequest(gomock.Any(), "/scan_result",
		gomock.Any(),
		gomock.Any(),
	).Times(0)

	wg := startCustomWatcher(t, ctx, kubeClientset, 1, corev1.PodFailed)

	s.Run(ctx)
	podList, err := kubeClientset.CoreV1().Pods(gitlabAgentNamespace).List(ctx, metav1.ListOptions{})
	assert.NoError(t, err)
	// Failed Trivy Pods should be deleted
	assert.Equal(t, 0, len(podList.Items))

	wg.Wait()
}

func TestScanEmptyNamespace(t *testing.T) {
	ctrl := gomock.NewController(t)
	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
	defer cancel()

	kubeClientset := fake.NewSimpleClientset()
	// Ensure that no Trivy Pods are created when there are no namespaces to be scanned
	kubeClientset.PrependReactor("create", "*", func(action clienttesting.Action) (bool, runtime.Object, error) {
		t.Errorf("Create should not be called")
		return true, nil, nil
	})

	logParser := &CustomLogParser{}
	api := mock_modagent.NewMockApi(ctrl)
	s := scanJob{
		log:                       zaptest.NewLogger(t),
		api:                       api,
		kubeClientset:             kubeClientset,
		gitlabAgentNamespace:      gitlabAgentNamespace,
		gitlabAgentServiceAccount: gitlabAgentServiceAccount,
		agentID:                   5000000000,
		targetNamespaces:          []string{}, //No namespace targetted
		logParser:                 logParser,
		resourceRequirements:      &defaultResourceRequirements,
	}

	api.EXPECT().MakeGitLabRequest(gomock.Any(), "/",
		gomock.Any(),
		gomock.Any(),
	).Times(0)

	api.EXPECT().MakeGitLabRequest(gomock.Any(), "/scan_result",
		gomock.Any(),
		gomock.Any(),
	).Times(0)

	wg := startCustomWatcher(t, ctx, kubeClientset, 1, corev1.PodSucceeded)

	s.Run(ctx)

	wg.Wait()
}

// Actual logs taken from output of Trivy k8s
var podLogsWithReport = []byte(`{
	"ClusterName": "",
	"Findings": [
	  {
	  "Namespace": "default",
	  "Kind": "Pod",
	  "Name": "nginx-unaccessible-6dfd584649-twqzw",
	  "Error": "scan error: unable to initialize a scanner: unable to initialize a docker scanner: 4 errors occurred:\\n\\t* unable to inspect the image (localhost:1234/nginx:latest): Cannot connect to the Docker daemon at unix:///var/run/docker.sock. Is the docker daemon running?\\n\\t* unable to initialize Podman client: no podman socket found: stat podman/podman.sock: no such file or directory\\n\\t* containerd socket not found: /run/containerd/containerd.sock\\n\\t* Get \"https://localhost:1234/v2/\": dial tcp [::1]:1234: connect: connection refused; Get \"http://localhost:1234/v2/\": dial tcp [::1]:1234: connect: connection refused\\n\\n"
	  }
	]
	}`)

// Actual logs with error taken from output of Trivy k8s
var podLogsWithErrorAndReport = []byte(fmt.Sprintf(`text2 2023-04-18T06:51:05.356Z	ERROR	Error during vulnerabilities scan: scan error: unable to initialize a scanner: unable to initialize a docker scanner: 4 errors occurred:
	  * unable to inspect the image (localhost:1234/nginx:latest): Cannot connect to the Docker daemon at unix:///var/run/docker.sock. Is the docker daemon running?
	  * unable to initialize Podman client: no podman socket found: stat podman/podman.sock: no such file or directory
	  * containerd socket not found: /run/containerd/containerd.sock
	  * Get "https://localhost:1234/v2/": dial tcp [::1]:1234: connect: connection refused; Get "http://localhost:1234/v2/": dial tcp [::1]:1234: connect: connection refused


  %s`, string(podLogsWithReport)))

func TestParsePodLogsToReportWithNoError(t *testing.T) {
	logParser := &logParserImpl{}
	consolidatedReport, err := logParser.ParsePodLogsToReport(podLogsWithReport)
	assert.NoError(t, err)

	var referenceConsolidatedReport ConsolidatedReport
	err = json.Unmarshal(podLogsWithReport, &referenceConsolidatedReport)
	assert.NoError(t, err)

	assert.Equal(t, referenceConsolidatedReport, consolidatedReport)
}

func TestParsePodLogsToReportWithError(t *testing.T) {
	logParser := &logParserImpl{}
	consolidatedReport, err := logParser.ParsePodLogsToReport(podLogsWithErrorAndReport)
	assert.NoError(t, err)

	var referenceConsolidatedReport ConsolidatedReport
	err = json.Unmarshal(podLogsWithReport, &referenceConsolidatedReport)
	assert.NoError(t, err)

	assert.Equal(t, referenceConsolidatedReport, consolidatedReport)
}

var publishedDate = time.Date(2020, time.December, 10, 0, 0, 0, 0, time.UTC)
var lastModifiedDate = time.Date(2022, time.October, 29, 0, 0, 0, 0, time.UTC)
var sampleResult = []Result{
	{
		Target: "nginx:1.16 (debian 10.3)",
		Class:  "os-pkgs",
		Type:   "alpine",
		Vulnerabilities: []DetectedVulnerability{
			{
				VulnerabilityID:  "CVE-2020-27350",
				PkgName:          "apt",
				InstalledVersion: "1.8.2",
				FixedVersion:     "1.8.2.2",
				PrimaryURL:       "https://avd.aquasec.com/nvd/cve-2020-27350",
				Vulnerability: Vulnerability{
					Title:       "apt: integer overflows and underflows while parsing .deb packages",
					Description: "APT had several integer overflows and underflows while parsing .deb packages, aka GHSL-2020-168 GHSL-2020-169, in files apt-pkg/contrib/extracttar.cc, apt-pkg/deb/debfile.cc, and apt-pkg/contrib/arfile.cc. This issue affects: apt 1.2.32ubuntu0 versions prior to 1.2.32ubuntu0.2; 1.6.12ubuntu0 versions prior to 1.6.12ubuntu0.2; 2.0.2ubuntu0 versions prior to 2.0.2ubuntu0.2; 2.1.10ubuntu0 versions prior to 2.1.10ubuntu0.1;",
					Severity:    "MEDIUM",
					References: []string{
						"https://access.redhat.com/security/cve/CVE-2020-27350",
						"https://bugs.launchpad.net/bugs/1899193",
					},
					PublishedDate:    &publishedDate,
					LastModifiedDate: &lastModifiedDate,
				},
			},
		},
	},
}

func setOwnerFindings(ownerKind string) Resource {
	return Resource{
		Namespace: namespace1,
		Kind:      ownerKind,
		Name:      "Owner name",
		Results:   sampleResult,
	}
}

func TestExcludeControlledPodWorkloads(t *testing.T) {
	ctrl := gomock.NewController(t)
	ctx := context.Background()

	kubeClientset := fake.NewSimpleClientset()
	logParser := &CustomLogParser{}
	api := mock_modagent.NewMockApi(ctrl)
	s := scanJob{
		log:                       zaptest.NewLogger(t),
		api:                       api,
		kubeClientset:             kubeClientset,
		gitlabAgentNamespace:      gitlabAgentNamespace,
		gitlabAgentServiceAccount: gitlabAgentServiceAccount,
		agentID:                   5000000000,
		targetNamespaces:          []string{namespace1},
		logParser:                 logParser,
	}

	// Pod spec with OwnerReference to simulate a controlled Pod
	isController := true
	podSpec := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:            podName,
			Namespace:       namespace1,
			OwnerReferences: []metav1.OwnerReference{},
		},
		Spec: corev1.PodSpec{
			Containers: []corev1.Container{
				{
					Name:  "nginx-container",
					Image: "nginx:1.16",
				},
			},
		},
	}

	podFindings := Resource{
		Namespace: namespace1,
		Kind:      "Pod",
		Name:      podName,
		Results:   sampleResult,
	}

	testCases := []struct {
		hasOwner                  bool
		ownerGroupVersion         string
		ownerKind                 string
		consolidatedReport        ConsolidatedReport
		expectedQualifiedFindings []Resource
	}{
		{
			hasOwner:          true,
			ownerGroupVersion: "v1",
			ownerKind:         "ReplicationController",
			consolidatedReport: ConsolidatedReport{
				Findings: []Resource{
					setOwnerFindings("ReplicationController"),
				},
			},
			expectedQualifiedFindings: []Resource{
				setOwnerFindings("ReplicationController"),
			},
		},
		{
			hasOwner:          true,
			ownerGroupVersion: "apps/v1",
			ownerKind:         "ReplicaSet",
			consolidatedReport: ConsolidatedReport{
				Findings: []Resource{
					setOwnerFindings("ReplicaSet"),
				},
			},
			expectedQualifiedFindings: []Resource{
				setOwnerFindings("ReplicaSet"),
			},
		},
		{
			hasOwner:          true,
			ownerGroupVersion: "apps/v1",
			ownerKind:         "StatefulSet",
			consolidatedReport: ConsolidatedReport{
				Findings: []Resource{
					setOwnerFindings("StatefulSet"),
				},
			},
			expectedQualifiedFindings: []Resource{
				setOwnerFindings("StatefulSet"),
			},
		},
		{
			hasOwner:          true,
			ownerGroupVersion: "apps/v1",
			ownerKind:         "DaemonSet",
			consolidatedReport: ConsolidatedReport{
				Findings: []Resource{
					setOwnerFindings("DaemonSet"),
				},
			},
			expectedQualifiedFindings: []Resource{
				setOwnerFindings("DaemonSet"),
			},
		},
		{
			hasOwner:          true,
			ownerGroupVersion: "batch/v1",
			ownerKind:         "Job",
			consolidatedReport: ConsolidatedReport{
				Findings: []Resource{
					setOwnerFindings("Job"),
				},
			},
			expectedQualifiedFindings: []Resource{
				setOwnerFindings("Job"),
			},
		},
		{
			hasOwner:          true,
			ownerGroupVersion: "custom/v1",
			ownerKind:         "custom",
			consolidatedReport: ConsolidatedReport{
				Findings: []Resource{
					setOwnerFindings("custom"),
					podFindings,
				},
			},
			expectedQualifiedFindings: []Resource{
				setOwnerFindings("custom"),
				podFindings,
			},
		},
		{
			hasOwner:          false,
			ownerGroupVersion: "",
			ownerKind:         "",
			consolidatedReport: ConsolidatedReport{
				Findings: []Resource{
					podFindings,
				},
			},
			expectedQualifiedFindings: []Resource{
				podFindings,
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.ownerKind, func(t *testing.T) {
			if tc.hasOwner {
				podSpec.ObjectMeta.OwnerReferences = []metav1.OwnerReference{
					{
						APIVersion: tc.ownerGroupVersion,
						Kind:       tc.ownerKind,
						Name:       "Owner name",
						Controller: &isController,
					},
				}
			}

			// Create Pod
			_, err := kubeClientset.CoreV1().Pods(namespace1).Create(ctx, podSpec, metav1.CreateOptions{})
			assert.NoError(t, err)

			qualifiedFindings := s.excludeControlledPodWorkloads(ctx, s.log, &tc.consolidatedReport, namespace1)

			assert.Equal(t, tc.expectedQualifiedFindings, qualifiedFindings)

			// Delete pod
			err = s.kubeClientset.CoreV1().Pods(namespace1).Delete(ctx, podName, metav1.DeleteOptions{})
			assert.NoError(t, err)
		})
	}
}

func getUuidsFromOptions(opts []modagent.GitLabRequestOption) ([]string, error) {
	body, err := getBodyFromOptions(opts)
	if err != nil {
		return nil, err
	}

	payload := new(reqBody)
	if err := json.Unmarshal(body, payload); err != nil {
		return nil, err
	}

	return payload.UUIDs, nil
}

func getBodyFromOptions(opts []modagent.GitLabRequestOption) ([]byte, error) {
	config, err := modagent.ApplyRequestOptions(opts)
	if err != nil {
		return nil, err
	}

	return io.ReadAll(config.Body)
}

func mockVulnerabilityCreate(ctx context.Context, path string, opts ...modagent.GitLabRequestOption) (*modagent.GitLabResponse, error) {
	body, err := getBodyFromOptions(opts)
	if err != nil {
		return nil, err
	}

	payload := new(Payload)
	if err = json.Unmarshal(body, payload); err != nil {
		return nil, err
	}

	response := &uuidResponse{
		// Technically not a UUID, but for testing
		// purposes it works as a suitible replacement.
		UUID: payload.Vulnerability.ID(),
	}

	var responseBytes []byte
	responseBytes, err = json.Marshal(response)
	if err != nil {
		return nil, err
	}

	return &modagent.GitLabResponse{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(bytes.NewReader(responseBytes)),
	}, nil
}

type CustomLogParser struct {
	err error
}

// ParsePodLogsToReport stubs the report that is parsed from a Trivy Pod's logs
func (m *CustomLogParser) ParsePodLogsToReport(logs []byte) (ConsolidatedReport, error) {
	report := ConsolidatedReport{
		Findings: []Resource{
			{
				Namespace: "some-namespace",
				Kind:      "Deployment",
				Name:      "nginx",
				Results:   sampleResult,
			},
		},
	}
	return report, m.err
}

// customWatcher is used to stub the Pod's watch interface
type customWatcher struct {
	eventChan chan watch.Event
}

func (cw *customWatcher) Stop() {}

func (cw *customWatcher) ResultChan() <-chan watch.Event {
	return cw.eventChan

}

// startCustomWatcher creates a custom watcher and waits for the specified number of Pods to be created before changing the Pod status to the desired status.
// This will ensure that the watcher in scanJob.startPodScanForNamespace does not wait forever and can continue processing.
// As we are using kubernetes fake.Clientset, Status.Phase would be "" instead of PodPending.
// Status.Phase would not change even after the Pod status has been updated since the customWatcher doesn't update the kubernetes object.
// As such, the moment the status has been updated, this function will exit.
func startCustomWatcher(t *testing.T, ctx context.Context, kubeClientset *fake.Clientset, podsToCreate int, desiredPodStatus corev1.PodPhase) *sync.WaitGroup {
	customWatcher := &customWatcher{
		eventChan: make(chan watch.Event),
	}
	kubeClientset.PrependWatchReactor("pods", func(action clienttesting.Action) (handled bool, ret watch.Interface, err error) {
		return true, customWatcher, nil
	})

	// Use WaitGroup to wait for Pod status to change before exiting the test function
	var wg sync.WaitGroup
	wg.Add(1)
	go func(wg *sync.WaitGroup) {
		defer wg.Done()

		for {
			select {
			case <-time.After(10 * time.Millisecond):
				podList, err := kubeClientset.CoreV1().Pods(gitlabAgentNamespace).List(ctx, metav1.ListOptions{})
				assert.NoError(t, err)

				// Only process if number of expected Pods have been created
				if len(podList.Items) == podsToCreate {
					for i := range podList.Items {
						createdPod := &podList.Items[i]
						updatedPod := createdPod.DeepCopy()
						updatedPod.Status.Phase = desiredPodStatus
						customWatcher.eventChan <- watch.Event{Type: watch.Modified, Object: updatedPod}
					}
					close(customWatcher.eventChan)
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}(&wg)

	return &wg
}

func TestGetPodSpec(t *testing.T) {
	s := &scanJob{
		gitlabAgentNamespace:      gitlabAgentNamespace,
		gitlabAgentServiceAccount: gitlabAgentServiceAccount,
		resourceRequirements:      &defaultResourceRequirements,
	}
	workloadsToScan := "Pod,ReplicaSet,ReplicationController,StatefulSet,DaemonSet,CronJob,Job"

	expectedPodSpec := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      podName,
			Namespace: s.gitlabAgentNamespace,
		},
		Spec: corev1.PodSpec{
			ServiceAccountName: s.gitlabAgentServiceAccount,
			Containers: []corev1.Container{
				{
					Name: "trivy",
					Resources: corev1.ResourceRequirements{
						Limits: corev1.ResourceList{
							corev1.ResourceMemory: resource.MustParse(defaultTrivyResourceLimitsMemory),
							corev1.ResourceCPU:    resource.MustParse(defaultTrivyResourceLimitsCpu),
						},
						Requests: corev1.ResourceList{
							corev1.ResourceMemory: resource.MustParse(defaultTrivyResourceRequestsMemory),
							corev1.ResourceCPU:    resource.MustParse(defaultTrivyResourceRequestsCpu),
						},
					},
					Image: trivyImageRef,
					Command: []string{
						"trivy",
						"k8s",
						workloadsToScan,
						"-n",
						namespace1,
						"--no-progress",
						"--quiet",
						"--report=summary",
						"--scanners=vuln",
						"--db-repository",
						"registry.gitlab.com/gitlab-org/security-products/dependencies/trivy-db-glad",
						"--format",
						"json",
					},
				},
			},
			RestartPolicy: corev1.RestartPolicyNever,
		},
	}

	actualPodSpec, err := s.getPodSpecForTrivyScanner(podName, namespace1)
	if assert.NoError(t, err) {
		assert.Equal(t, expectedPodSpec, actualPodSpec)
	}
}

func TestGetPodSpecWithValidationError(t *testing.T) {
	testCases := []struct {
		description         string
		inputLimitsCpu      string
		inputLimitsMemory   string
		inputRequestsCpu    string
		inputRequestsMemory string
	}{
		{
			description:         "invalid limits cpu",
			inputLimitsCpu:      "a",
			inputLimitsMemory:   "",
			inputRequestsCpu:    "",
			inputRequestsMemory: "",
		},
		{
			description:         "invalid limits memory",
			inputLimitsCpu:      "",
			inputLimitsMemory:   "a",
			inputRequestsCpu:    "",
			inputRequestsMemory: "",
		},
		{
			description:         "invalid requests cpu",
			inputLimitsCpu:      "",
			inputLimitsMemory:   "",
			inputRequestsCpu:    "a",
			inputRequestsMemory: "",
		},
		{
			description:         "invalid requests memory",
			inputLimitsCpu:      "",
			inputLimitsMemory:   "",
			inputRequestsCpu:    "",
			inputRequestsMemory: "a",
		},
	}
	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			s := &scanJob{
				gitlabAgentNamespace:      gitlabAgentNamespace,
				gitlabAgentServiceAccount: gitlabAgentServiceAccount,
				resourceRequirements: &agentcfg.ResourceRequirements{
					Limits: &agentcfg.Resource{
						Cpu:    tc.inputLimitsCpu,
						Memory: tc.inputLimitsMemory,
					},
					Requests: &agentcfg.Resource{
						Cpu:    tc.inputRequestsCpu,
						Memory: tc.inputRequestsMemory,
					},
				},
			}

			_, err := s.getPodSpecForTrivyScanner(podName, namespace1)
			assert.ErrorIs(t, err, resource.ErrFormatWrong)
		})
	}
}