// Copyright 2015 clair authors
//
// 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 database

import (
	"testing"

	cerrors "github.com/coreos/clair/utils/errors"
	"github.com/coreos/clair/utils/types"
	"github.com/stretchr/testify/assert"
)

func TestVulnerability(t *testing.T) {
	Open("memstore", "")
	defer Close()

	// Insert invalid vulnerabilities
	for _, vulnerability := range []Vulnerability{
		Vulnerability{ID: "", Link: "link1", Priority: types.Medium, FixedInNodes: []string{"pkg1"}},
		Vulnerability{ID: "test1", Link: "", Priority: types.Medium, FixedInNodes: []string{"pkg1"}},
		Vulnerability{ID: "test1", Link: "link1", Priority: "InvalidPriority", FixedInNodes: []string{"pkg1"}},
		Vulnerability{ID: "test1", Link: "link1", Priority: types.Medium, FixedInNodes: []string{}},
	} {
		_, err := InsertVulnerabilities([]*Vulnerability{&vulnerability})
		assert.Error(t, err)
	}

	// Some data
	vuln1 := &Vulnerability{ID: "test1", Link: "link1", Priority: types.Medium, Description: "testDescription1", FixedInNodes: []string{"pkg1"}}
	vuln2 := &Vulnerability{ID: "test2", Link: "link2", Priority: types.High, Description: "testDescription2", FixedInNodes: []string{"pkg1", "pkg2"}}
	vuln3 := &Vulnerability{ID: "test3", Link: "link3", Priority: types.High, FixedInNodes: []string{"pkg3"}} // Empty description

	// Insert some vulnerabilities
	_, err := InsertVulnerabilities([]*Vulnerability{vuln1, vuln2, vuln3})
	if assert.Nil(t, err) {
		// Find one of the vulnerabilities we just inserted and verify its content
		v1, err := FindOneVulnerability(vuln1.ID, FieldVulnerabilityAll)
		if assert.Nil(t, err) && assert.NotNil(t, v1) {
			assert.Equal(t, vuln1.ID, v1.ID)
			assert.Equal(t, vuln1.Link, v1.Link)
			assert.Equal(t, vuln1.Priority, v1.Priority)
			assert.Equal(t, vuln1.Description, v1.Description)
			if assert.Len(t, v1.FixedInNodes, 1) {
				assert.Equal(t, vuln1.FixedInNodes[0], v1.FixedInNodes[0])
			}
		}

		// Ensure that vulnerabilities with empty descriptions work as well
		v3, err := FindOneVulnerability(vuln3.ID, FieldVulnerabilityAll)
		if assert.Nil(t, err) && assert.NotNil(t, v3) {
			assert.Equal(t, vuln3.Description, v3.Description)
		}

		// Find vulnerabilities by fixed packages
		vulnsFixedInPkg1AndPkg3, err := FindAllVulnerabilitiesByFixedIn([]string{"pkg2", "pkg3"}, FieldVulnerabilityAll)
		assert.Nil(t, err)
		assert.Len(t, vulnsFixedInPkg1AndPkg3, 2)

		// Delete vulnerability
		if assert.Nil(t, DeleteVulnerability(vuln1.ID)) {
			v1, err := FindOneVulnerability(vuln1.ID, FieldVulnerabilityAll)
			assert.Equal(t, cerrors.ErrNotFound, err)
			assert.Nil(t, v1)
		}
	}

	// Update a vulnerability and verify its new content
	pkg1 := &Package{OS: "testOS", Name: "testpkg1", Version: types.NewVersionUnsafe("1.0")}
	InsertPackages([]*Package{pkg1})
	vuln5 := &Vulnerability{ID: "test5", Link: "link5", Priority: types.Medium, Description: "testDescription5", FixedInNodes: []string{pkg1.Node}}

	_, err = InsertVulnerabilities([]*Vulnerability{vuln5})
	if assert.Nil(t, err) {
		// Partial updates
		// # Just a field update
		vuln5b := &Vulnerability{ID: "test5", Priority: types.High}
		_, err := InsertVulnerabilities([]*Vulnerability{vuln5b})
		if assert.Nil(t, err) {
			v5b, err := FindOneVulnerability(vuln5b.ID, FieldVulnerabilityAll)
			if assert.Nil(t, err) && assert.NotNil(t, v5b) {
				assert.Equal(t, vuln5b.ID, v5b.ID)
				assert.Equal(t, vuln5b.Priority, v5b.Priority)

				if assert.Len(t, v5b.FixedInNodes, 1) {
					assert.Contains(t, v5b.FixedInNodes, pkg1.Node)
				}
			}
		}

		// # Just a field update, twice in the same transaction
		vuln5b1 := &Vulnerability{ID: "test5", Link: "http://foo.bar"}
		vuln5b2 := &Vulnerability{ID: "test5", Link: "http://bar.foo"}
		_, err = InsertVulnerabilities([]*Vulnerability{vuln5b1, vuln5b2})
		if assert.Nil(t, err) {
			v5b2, err := FindOneVulnerability(vuln5b2.ID, FieldVulnerabilityAll)
			if assert.Nil(t, err) && assert.NotNil(t, v5b2) {
				assert.Equal(t, vuln5b2.Link, v5b2.Link)
			}
		}

		// # All fields except fixedIn update
		vuln5c := &Vulnerability{ID: "test5", Link: "link5c", Priority: types.Critical, Description: "testDescription5c"}
		_, err = InsertVulnerabilities([]*Vulnerability{vuln5c})
		if assert.Nil(t, err) {
			v5c, err := FindOneVulnerability(vuln5c.ID, FieldVulnerabilityAll)
			if assert.Nil(t, err) && assert.NotNil(t, v5c) {
				assert.Equal(t, vuln5c.ID, v5c.ID)
				assert.Equal(t, vuln5c.Link, v5c.Link)
				assert.Equal(t, vuln5c.Priority, v5c.Priority)
				assert.Equal(t, vuln5c.Description, v5c.Description)

				if assert.Len(t, v5c.FixedInNodes, 1) {
					assert.Contains(t, v5c.FixedInNodes, pkg1.Node)
				}
			}
		}

		// Complete update
		pkg2 := &Package{OS: "testOS", Name: "testpkg1", Version: types.NewVersionUnsafe("1.1")}
		pkg3 := &Package{OS: "testOS", Name: "testpkg2", Version: types.NewVersionUnsafe("1.0")}
		InsertPackages([]*Package{pkg2, pkg3})
		vuln5d := &Vulnerability{ID: "test5", Link: "link5d", Priority: types.Low, Description: "testDescription5d", FixedInNodes: []string{pkg2.Node, pkg3.Node}}

		_, err = InsertVulnerabilities([]*Vulnerability{vuln5d})
		if assert.Nil(t, err) {
			v5d, err := FindOneVulnerability(vuln5d.ID, FieldVulnerabilityAll)
			if assert.Nil(t, err) && assert.NotNil(t, v5d) {
				assert.Equal(t, vuln5d.ID, v5d.ID)
				assert.Equal(t, vuln5d.Link, v5d.Link)
				assert.Equal(t, vuln5d.Priority, v5d.Priority)
				assert.Equal(t, vuln5d.Description, v5d.Description)

				// Here, we ensure that a vulnerability can only be fixed by one package of a given branch at a given time
				// And that we can add new fixed packages as well
				if assert.Len(t, v5d.FixedInNodes, 2) {
					assert.NotContains(t, v5d.FixedInNodes, pkg1.Node)
				}
			}
		}
	}

	// Create and update a vulnerability's packages (and from the same branch) in the same batch
	pkg1 = &Package{OS: "testOS", Name: "testpkg1", Version: types.NewVersionUnsafe("1.0")}
	pkg1b := &Package{OS: "testOS", Name: "testpkg1", Version: types.NewVersionUnsafe("1.1")}
	InsertPackages([]*Package{pkg1, pkg1b})
	// # A vulnerability can't be inserted if fixed by two packages of the same branch
	_, err = InsertVulnerabilities([]*Vulnerability{&Vulnerability{ID: "test6", Link: "link6", Priority: types.Medium, Description: "testDescription6", FixedInNodes: []string{pkg1.Node, pkg1b.Node}}})
	assert.Error(t, err)
	// # Two updates of the same vulnerability in the same batch with packages of the same branch
	pkg0 := &Package{OS: "testOS", Name: "testpkg0", Version: types.NewVersionUnsafe("1.0")}
	InsertPackages([]*Package{pkg0})
	_, err = InsertVulnerabilities([]*Vulnerability{&Vulnerability{ID: "test7", Link: "link7", Priority: types.Medium, Description: "testDescription7", FixedInNodes: []string{pkg0.Node}}})
	if assert.Nil(t, err) {
		vuln7b := &Vulnerability{ID: "test7", FixedInNodes: []string{pkg1.Node}}
		vuln7c := &Vulnerability{ID: "test7", FixedInNodes: []string{pkg1b.Node}}
		_, err = InsertVulnerabilities([]*Vulnerability{vuln7b, vuln7c})
		if assert.Nil(t, err) {
			v7, err := FindOneVulnerability("test7", FieldVulnerabilityAll)
			if assert.Nil(t, err) && assert.Len(t, v7.FixedInNodes, 2) {
				assert.Contains(t, v7.FixedInNodes, pkg0.Node)
				assert.NotContains(t, v7.FixedInNodes, pkg1.Node)
				assert.Contains(t, v7.FixedInNodes, pkg1b.Node)
			}

			// # A vulnerability can't be updated if fixed by two packages of the same branch
			_, err = InsertVulnerabilities([]*Vulnerability{&Vulnerability{ID: "test7", FixedInNodes: []string{pkg1.Node, pkg1b.Node}}})
			assert.Error(t, err)
		}
	}
}

func TestInsertVulnerabilityNotifications(t *testing.T) {
	Open("memstore", "")
	defer Close()

	pkg1 := &Package{OS: "testOS", Name: "testpkg1", Version: types.NewVersionUnsafe("1.0")}
	pkg1b := &Package{OS: "testOS", Name: "testpkg1", Version: types.NewVersionUnsafe("1.2")}
	pkg2 := &Package{OS: "testOS", Name: "testpkg2", Version: types.NewVersionUnsafe("1.0")}
	InsertPackages([]*Package{pkg1, pkg1b, pkg2})

	// NewVulnerabilityNotification
	vuln1 := &Vulnerability{ID: "test1", Link: "link1", Priority: types.Medium, Description: "testDescription1", FixedInNodes: []string{pkg1.Node}}
	vuln2 := &Vulnerability{ID: "test2", Link: "link2", Priority: types.High, Description: "testDescription2", FixedInNodes: []string{pkg1.Node, pkg2.Node}}
	vuln1b := &Vulnerability{ID: "test1", Priority: types.High, FixedInNodes: []string{"pkg3"}}
	notifications, err := InsertVulnerabilities([]*Vulnerability{vuln1, vuln2, vuln1b})
	if assert.Nil(t, err) {
		// We should only have two NewVulnerabilityNotification notifications: one for test1 and one for test2
		// We should not have a VulnerabilityPriorityIncreasedNotification or a VulnerabilityPackageChangedNotification
		// for test1 because it is in the same batch
		if assert.Len(t, notifications, 2) {
			for _, n := range notifications {
				_, ok := n.(*NewVulnerabilityNotification)
				assert.True(t, ok)
			}
		}
	}

	// VulnerabilityPriorityIncreasedNotification
	vuln1c := &Vulnerability{ID: "test1", Priority: types.Critical}
	notifications, err = InsertVulnerabilities([]*Vulnerability{vuln1c})
	if assert.Nil(t, err) {
		if assert.Len(t, notifications, 1) {
			if nn, ok := notifications[0].(*VulnerabilityPriorityIncreasedNotification); assert.True(t, ok) {
				assert.Equal(t, vuln1b.Priority, nn.OldPriority)
				assert.Equal(t, vuln1c.Priority, nn.NewPriority)
			}
		}
	}

	notifications, err = InsertVulnerabilities([]*Vulnerability{&Vulnerability{ID: "test1", Priority: types.Low}})
	assert.Nil(t, err)
	assert.Len(t, notifications, 0)

	// VulnerabilityPackageChangedNotification
	vuln1e := &Vulnerability{ID: "test1", FixedInNodes: []string{pkg1b.Node}}
	vuln1f := &Vulnerability{ID: "test1", FixedInNodes: []string{pkg2.Node}}
	notifications, err = InsertVulnerabilities([]*Vulnerability{vuln1e, vuln1f})
	if assert.Nil(t, err) {
		if assert.Len(t, notifications, 1) {
			if nn, ok := notifications[0].(*VulnerabilityPackageChangedNotification); assert.True(t, ok) {
				// Here, we say that pkg1b fixes the vulnerability, but as pkg1b is in
				// the same branch as pkg1, pkg1 should be removed and pkg1b added
				// We also add pkg2 as fixed
				assert.Contains(t, nn.AddedFixedInNodes, pkg1b.Node)
				assert.Contains(t, nn.RemovedFixedInNodes, pkg1.Node)

				assert.Contains(t, nn.AddedFixedInNodes, pkg2.Node)
			}
		}
	}
}