// Copyright (c) 2018 Intel Corporation
//
// SPDX-License-Identifier: Apache-2.0
//

package main

import (
	"bytes"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"testing"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/urfave/cli"
)

func setupCheckHostIsVMContainerCapable(assert *assert.Assertions, cpuInfoFile string, cpuData []testCPUData, moduleData []testModuleData) {
	createModules(assert, cpuInfoFile, moduleData)

	// all the modules files have now been created, so deal with the
	// cpuinfo data.
	for _, d := range cpuData {
		err := makeCPUInfoFile(cpuInfoFile, d.vendorID, d.flags)
		assert.NoError(err)

		details := vmContainerCapableDetails{
			cpuInfoFile:           cpuInfoFile,
			requiredCPUFlags:      archRequiredCPUFlags,
			requiredCPUAttribs:    archRequiredCPUAttribs,
			requiredKernelModules: archRequiredKernelModules,
		}

		err = hostIsVMContainerCapable(details)
		if d.expectError {
			assert.Error(err)
		} else {
			assert.NoError(err)
		}
	}
}

func TestCCCheckCLIFunction(t *testing.T) {
	assert := assert.New(t)

	dir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	savedSysModuleDir := sysModuleDir
	savedProcCPUInfo := procCPUInfo

	cpuInfoFile := filepath.Join(dir, "cpuinfo")

	// XXX: override
	sysModuleDir = filepath.Join(dir, "sys/module")
	procCPUInfo = cpuInfoFile

	defer func() {
		sysModuleDir = savedSysModuleDir
		procCPUInfo = savedProcCPUInfo
	}()

	err = os.MkdirAll(sysModuleDir, testDirMode)
	if err != nil {
		t.Fatal(err)
	}

	var cpuData []testCPUData
	var moduleData []testModuleData

	if cpuType == cpuTypeIntel {
		cpuData = []testCPUData{
			{archGenuineIntel, "lm vmx sse4_1", false},
		}

		moduleData = []testModuleData{
			{filepath.Join(sysModuleDir, "kvm_intel/parameters/unrestricted_guest"), false, "Y"},
			{filepath.Join(sysModuleDir, "kvm_intel/parameters/nested"), false, "Y"},
		}
	} else if cpuType == cpuTypeAMD {
		cpuData = []testCPUData{
			{archAuthenticAMD, "lm svm sse4_1", false},
		}

		moduleData = []testModuleData{
			{filepath.Join(sysModuleDir, "kvm_amd/parameters/nested"), false, "1"},
		}
	}

	devNull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0666)
	assert.NoError(err)
	defer devNull.Close()

	savedLogOutput := kataLog.Logger.Out

	// discard normal output
	kataLog.Logger.Out = devNull

	defer func() {
		kataLog.Logger.Out = savedLogOutput
	}()

	setupCheckHostIsVMContainerCapable(assert, cpuInfoFile, cpuData, moduleData)

	ctx := createCLIContext(nil)
	ctx.App.Name = "foo"

	// create buffer to save logger output
	buf := &bytes.Buffer{}

	// capture output this time
	kataLog.Logger.Out = buf

	fn, ok := kataCheckCLICommand.Action.(func(context *cli.Context) error)
	assert.True(ok)

	err = fn(ctx)
	assert.NoError(err)

	output := buf.String()

	for _, c := range cpuData {
		assert.True(findAnchoredString(output, c.vendorID))
		for _, flag := range strings.Fields(c.flags) {
			assert.True(findAnchoredString(output, flag))
		}
	}

	for _, m := range moduleData {
		name := path.Base(m.path)
		assert.True(findAnchoredString(output, name))
	}
}

func TestCheckCheckKernelModulesNoNesting(t *testing.T) {
	assert := assert.New(t)

	dir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	savedSysModuleDir := sysModuleDir
	savedProcCPUInfo := procCPUInfo

	cpuInfoFile := filepath.Join(dir, "cpuinfo")

	// XXX: override
	sysModuleDir = filepath.Join(dir, "sys/module")
	procCPUInfo = cpuInfoFile

	defer func() {
		sysModuleDir = savedSysModuleDir
		procCPUInfo = savedProcCPUInfo
	}()

	err = os.MkdirAll(sysModuleDir, testDirMode)
	if err != nil {
		t.Fatal(err)
	}

	requiredModules := map[string]kernelModule{
		"kvm_intel": {
			desc: "Intel KVM",
			parameters: map[string]string{
				"nested":             "Y",
				"unrestricted_guest": "Y",
			},
		},
	}

	actualModuleData := []testModuleData{
		{filepath.Join(sysModuleDir, "kvm"), true, ""},
		{filepath.Join(sysModuleDir, "kvm_intel"), true, ""},
		{filepath.Join(sysModuleDir, "kvm_intel/parameters/unrestricted_guest"), false, "Y"},

		// XXX: force a warning
		{filepath.Join(sysModuleDir, "kvm_intel/parameters/nested"), false, "N"},
	}

	vendor := archGenuineIntel
	flags := "vmx lm sse4_1 hypervisor"

	_, err = checkKernelModules(requiredModules, archKernelParamHandler)
	// no cpuInfoFile yet
	assert.Error(err)

	createModules(assert, cpuInfoFile, actualModuleData)

	err = makeCPUInfoFile(cpuInfoFile, vendor, flags)
	assert.NoError(err)

	count, err := checkKernelModules(requiredModules, archKernelParamHandler)
	assert.NoError(err)
	assert.Equal(count, uint32(0))

	// create buffer to save logger output
	buf := &bytes.Buffer{}

	savedLogOutput := kataLog.Logger.Out

	defer func() {
		kataLog.Logger.Out = savedLogOutput
	}()

	kataLog.Logger.Out = buf

	count, err = checkKernelModules(requiredModules, archKernelParamHandler)

	assert.NoError(err)
	assert.Equal(count, uint32(0))

	re := regexp.MustCompile(`\bwarning\b.*\bnested\b`)
	matches := re.FindAllStringSubmatch(buf.String(), -1)
	assert.NotEmpty(matches)
}

func TestCheckCheckKernelModulesNoUnrestrictedGuest(t *testing.T) {
	assert := assert.New(t)

	dir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	savedSysModuleDir := sysModuleDir
	savedProcCPUInfo := procCPUInfo

	cpuInfoFile := filepath.Join(dir, "cpuinfo")

	// XXX: override
	sysModuleDir = filepath.Join(dir, "sys/module")
	procCPUInfo = cpuInfoFile

	defer func() {
		sysModuleDir = savedSysModuleDir
		procCPUInfo = savedProcCPUInfo
	}()

	err = os.MkdirAll(sysModuleDir, testDirMode)
	if err != nil {
		t.Fatal(err)
	}

	requiredModules := map[string]kernelModule{
		"kvm_intel": {
			desc: "Intel KVM",
			parameters: map[string]string{
				"nested":             "Y",
				"unrestricted_guest": "Y",
			},
		},
	}

	actualModuleData := []testModuleData{
		{filepath.Join(sysModuleDir, "kvm"), true, ""},
		{filepath.Join(sysModuleDir, "kvm_intel"), true, ""},
		{filepath.Join(sysModuleDir, "kvm_intel/parameters/nested"), false, "Y"},

		// XXX: force a failure on non-VMM systems
		{filepath.Join(sysModuleDir, "kvm_intel/parameters/unrestricted_guest"), false, "N"},
	}

	vendor := archGenuineIntel
	flags := "vmx lm sse4_1"

	_, err = checkKernelModules(requiredModules, archKernelParamHandler)
	// no cpuInfoFile yet
	assert.Error(err)

	err = makeCPUInfoFile(cpuInfoFile, vendor, flags)
	assert.NoError(err)

	createModules(assert, cpuInfoFile, actualModuleData)

	count, err := checkKernelModules(requiredModules, archKernelParamHandler)

	assert.NoError(err)
	// fails due to unrestricted_guest not being available
	assert.Equal(count, uint32(1))

	// pretend test is running under a hypervisor
	flags += " hypervisor"

	// recreate
	err = makeCPUInfoFile(cpuInfoFile, vendor, flags)
	assert.NoError(err)

	// create buffer to save logger output
	buf := &bytes.Buffer{}

	savedLogOutput := kataLog.Logger.Out

	defer func() {
		kataLog.Logger.Out = savedLogOutput
	}()

	kataLog.Logger.Out = buf

	count, err = checkKernelModules(requiredModules, archKernelParamHandler)

	// no error now because running under a hypervisor
	assert.NoError(err)
	assert.Equal(count, uint32(0))

	re := regexp.MustCompile(`\bwarning\b.*\bunrestricted_guest\b`)
	matches := re.FindAllStringSubmatch(buf.String(), -1)
	assert.NotEmpty(matches)
}

func TestCheckHostIsVMContainerCapable(t *testing.T) {
	assert := assert.New(t)

	dir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	savedSysModuleDir := sysModuleDir
	savedProcCPUInfo := procCPUInfo

	cpuInfoFile := filepath.Join(dir, "cpuinfo")

	// XXX: override
	sysModuleDir = filepath.Join(dir, "sys/module")
	procCPUInfo = cpuInfoFile

	defer func() {
		sysModuleDir = savedSysModuleDir
		procCPUInfo = savedProcCPUInfo
	}()

	err = os.MkdirAll(sysModuleDir, testDirMode)
	if err != nil {
		t.Fatal(err)
	}

	var cpuData []testCPUData
	var moduleData []testModuleData

	if cpuType == cpuTypeIntel {
		cpuData = []testCPUData{
			{"", "", true},
			{"Intel", "", true},
			{archGenuineIntel, "", true},
			{archGenuineIntel, "lm", true},
			{archGenuineIntel, "lm vmx", true},
			{archGenuineIntel, "lm vmx sse4_1", false},
		}

		moduleData = []testModuleData{
			{filepath.Join(sysModuleDir, "kvm"), true, ""},
			{filepath.Join(sysModuleDir, "kvm_intel"), true, ""},
			{filepath.Join(sysModuleDir, "kvm_intel/parameters/nested"), false, "Y"},
			{filepath.Join(sysModuleDir, "kvm_intel/parameters/unrestricted_guest"), false, "Y"},
		}
	} else if cpuType == cpuTypeAMD {
		cpuData = []testCPUData{
			{"", "", true},
			{"AMD", "", true},
			{archAuthenticAMD, "", true},
			{archAuthenticAMD, "lm", true},
			{archAuthenticAMD, "lm svm", true},
			{archAuthenticAMD, "lm svm sse4_1", false},
		}

		moduleData = []testModuleData{
			{filepath.Join(sysModuleDir, "kvm"), true, ""},
			{filepath.Join(sysModuleDir, "kvm_amd"), true, ""},
			{filepath.Join(sysModuleDir, "kvm_amd/parameters/nested"), false, "1"},
		}
	}

	setupCheckHostIsVMContainerCapable(assert, cpuInfoFile, cpuData, moduleData)

	// remove the modules to force a failure
	err = os.RemoveAll(sysModuleDir)
	assert.NoError(err)

	details := vmContainerCapableDetails{
		cpuInfoFile:           cpuInfoFile,
		requiredCPUFlags:      archRequiredCPUFlags,
		requiredCPUAttribs:    archRequiredCPUAttribs,
		requiredKernelModules: archRequiredKernelModules,
	}

	err = hostIsVMContainerCapable(details)
	assert.Error(err)
}

func TestArchKernelParamHandler(t *testing.T) {
	assert := assert.New(t)

	type testData struct {
		onVMM        bool
		fields       logrus.Fields
		msg          string
		expectIgnore bool
	}

	data := []testData{
		{true, logrus.Fields{}, "", false},
		{false, logrus.Fields{}, "", false},

		{
			false,
			logrus.Fields{
				// wrong type
				"parameter": 123,
			},
			"foo",
			false,
		},

		{
			false,
			logrus.Fields{
				"parameter": "unrestricted_guest",
			},
			"",
			false,
		},

		{
			true,
			logrus.Fields{
				"parameter": "unrestricted_guest",
			},
			"",
			true,
		},

		{
			false,
			logrus.Fields{
				"parameter": "nested",
			},
			"",
			true,
		},
	}

	for i, d := range data {
		result := archKernelParamHandler(d.onVMM, d.fields, d.msg)
		if d.expectIgnore {
			assert.True(result, "test %d (%+v)", i, d)
		} else {
			assert.False(result, "test %d (%+v)", i, d)
		}
	}
}

func TestKvmIsUsable(t *testing.T) {
	assert := assert.New(t)

	dir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	savedKvmDevice := kvmDevice
	fakeKVMDevice := filepath.Join(dir, "kvm")
	kvmDevice = fakeKVMDevice

	defer func() {
		kvmDevice = savedKvmDevice
	}()

	err = kvmIsUsable()
	assert.Error(err)

	err = createEmptyFile(fakeKVMDevice)
	assert.NoError(err)

	err = kvmIsUsable()
	assert.Error(err)
}

func TestGetCPUDetails(t *testing.T) {
	genericTestGetCPUDetails(t)
}

func TestSetCPUtype(t *testing.T) {
	assert := assert.New(t)

	savedArchRequiredCPUFlags := archRequiredCPUFlags
	savedArchRequiredCPUAttribs := archRequiredCPUAttribs
	savedArchRequiredKernelModules := archRequiredKernelModules

	defer func() {
		archRequiredCPUFlags = savedArchRequiredCPUFlags
		archRequiredCPUAttribs = savedArchRequiredCPUAttribs
		archRequiredKernelModules = savedArchRequiredKernelModules
	}()

	archRequiredCPUFlags = map[string]string{}
	archRequiredCPUAttribs = map[string]string{}
	archRequiredKernelModules = map[string]kernelModule{}

	setCPUtype()

	assert.NotEmpty(archRequiredCPUFlags)
	assert.NotEmpty(archRequiredCPUAttribs)
	assert.NotEmpty(archRequiredKernelModules)

	assert.Equal(archRequiredCPUFlags["vmx"], "Virtualization support")

	_, ok := archRequiredKernelModules["kvm"]
	assert.True(ok)
}