// -*- Mode: Go; indent-tabs-mode: t -*-
//go:build !nosecboot

/*
 * Copyright (C) 2025 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package secboot_test

import (
	"context"
	"encoding/json"
	"errors"

	"github.com/canonical/go-tpm2"
	sb_efi "github.com/snapcore/secboot/efi"
	sb_preinstall "github.com/snapcore/secboot/efi/preinstall"
	. "gopkg.in/check.v1"

	"github.com/snapcore/snapd/logger"
	"github.com/snapcore/snapd/secboot"
	"github.com/snapcore/snapd/testutil"
)

type preinstallSuite struct {
	testutil.BaseTest
}

var _ = Suite(&preinstallSuite{})

func (s *preinstallSuite) SetUpTest(c *C) {
}

// CompoundPreinstallCheckError implements sb_preinstall.CompoundError and is
// used to mimic compound errors that would normally be returned by secboot.
type CompoundPreinstallCheckError struct {
	Errs []error
}

func (e *CompoundPreinstallCheckError) Error() string {
	return "n/a"
}

func (e *CompoundPreinstallCheckError) Unwrap() []error {
	return e.Errs
}

func (s *preinstallSuite) TestConvertPreinstallCheckErrorActions(c *C) {
	testCases := []struct {
		actions  []sb_preinstall.Action
		expected []string
	}{
		{nil, nil},
		{[]sb_preinstall.Action{}, []string{}},
		{[]sb_preinstall.Action{sb_preinstall.ActionReboot, sb_preinstall.ActionShutdown}, []string{"reboot", "shutdown"}},
	}

	for _, tc := range testCases {
		convertedActions := secboot.ConvertPreinstallCheckErrorActions(tc.actions)
		c.Check(convertedActions, DeepEquals, tc.expected)
	}
}

func (s *preinstallSuite) TestConvertPreinstallCheckErrorType(c *C) {
	kindAndActionsErr := sb_preinstall.NewWithKindAndActionsError(
		sb_preinstall.ErrorKindTPMHierarchiesOwned,
		sb_preinstall.TPM2OwnedHierarchiesError{WithAuthValue: tpm2.HandleList{tpm2.HandleLockout}, WithAuthPolicy: tpm2.HandleList{tpm2.HandleOwner}},
		[]sb_preinstall.Action{sb_preinstall.ActionRebootToFWSettings},
		errors.New("error with TPM2 device: one or more of the TPM hierarchies is already owned"),
	)

	var errorDetails secboot.PreinstallErrorDetails
	c.Assert(func() { errorDetails = secboot.ConvertPreinstallCheckErrorType(nil) }, PanicMatches, "runtime error: invalid memory address or nil pointer dereference")
	errorDetails = secboot.ConvertPreinstallCheckErrorType(kindAndActionsErr)
	c.Assert(errorDetails, DeepEquals, secboot.PreinstallErrorDetails{
		Kind:    "tpm-hierarchies-owned",
		Message: "error with TPM2 device: one or more of the TPM hierarchies is already owned",
		Args: map[string]json.RawMessage{
			"with-auth-value":  json.RawMessage(`[1073741834]`),
			"with-auth-policy": json.RawMessage(`[1073741825]`),
		},
		Actions: []string{"reboot-to-fw-settings"},
	})
}

func (s *preinstallSuite) TestUnwrapPreinstallCheckErrorCompound(c *C) {
	compoundError := &CompoundPreinstallCheckError{
		[]error{
			sb_preinstall.NewWithKindAndActionsError(
				sb_preinstall.ErrorKindTPMHierarchiesOwned,
				sb_preinstall.TPM2OwnedHierarchiesError{WithAuthValue: tpm2.HandleList{tpm2.HandleLockout}},
				[]sb_preinstall.Action{sb_preinstall.ActionRebootToFWSettings},
				errors.New("error with TPM2 device: one or more of the TPM hierarchies is already owned"),
			),
			sb_preinstall.NewWithKindAndActionsError(
				sb_preinstall.ErrorKindTPMDeviceLockout,
				&sb_preinstall.TPMDeviceLockoutArgs{IntervalDuration: 7200000000000, TotalDuration: 230400000000000},
				[]sb_preinstall.Action{sb_preinstall.ActionRebootToFWSettings},
				errors.New("error with TPM2 device: TPM is in DA lockout mode"),
			),
			sb_preinstall.NewWithKindAndActionsError(
				sb_preinstall.ErrorKindNone,
				nil,
				nil,
				nil,
			),
		},
	}

	errorDetails, err := secboot.UnwrapPreinstallCheckError(compoundError)
	c.Assert(err, IsNil)
	c.Assert(errorDetails, DeepEquals, []secboot.PreinstallErrorDetails{
		{
			Kind:    "tpm-hierarchies-owned",
			Message: "error with TPM2 device: one or more of the TPM hierarchies is already owned",
			Args: map[string]json.RawMessage{
				"with-auth-value":  json.RawMessage(`[1073741834]`),
				"with-auth-policy": json.RawMessage(`null`),
			},
			Actions: []string{"reboot-to-fw-settings"},
		},
		{
			Kind:    "tpm-device-lockout",
			Message: "error with TPM2 device: TPM is in DA lockout mode",
			Args: map[string]json.RawMessage{
				"interval-duration": json.RawMessage(`7200000000000`),
				"total-duration":    json.RawMessage(`230400000000000`),
			},
			Actions: []string{"reboot-to-fw-settings"},
		},
		{
			Kind:    "",
			Message: "<nil>",
			Args:    nil,
			Actions: nil,
		},
	})

	jsn, err := json.MarshalIndent(errorDetails, "", "  ")
	c.Assert(err, IsNil)
	const expectedJson = `[
  {
    "kind": "tpm-hierarchies-owned",
    "message": "error with TPM2 device: one or more of the TPM hierarchies is already owned",
    "args": {
      "with-auth-policy": null,
      "with-auth-value": [
        1073741834
      ]
    },
    "actions": [
      "reboot-to-fw-settings"
    ]
  },
  {
    "kind": "tpm-device-lockout",
    "message": "error with TPM2 device: TPM is in DA lockout mode",
    "args": {
      "interval-duration": 7200000000000,
      "total-duration": 230400000000000
    },
    "actions": [
      "reboot-to-fw-settings"
    ]
  },
  {
    "kind": "",
    "message": "\u003cnil\u003e"
  }
]`
	c.Assert(string(jsn), Equals, expectedJson)
}

func (s *preinstallSuite) TestUnwrapPreinstallCheckErrorFailCompoundUnexpectedType(c *C) {
	compoundError := &CompoundPreinstallCheckError{
		[]error{
			sb_preinstall.NewWithKindAndActionsError(
				sb_preinstall.ErrorKindTPMHierarchiesOwned,
				sb_preinstall.TPM2OwnedHierarchiesError{WithAuthValue: tpm2.HandleList{tpm2.HandleLockout}},
				[]sb_preinstall.Action{sb_preinstall.ActionRebootToFWSettings},
				errors.New("error with TPM2 device: one or more of the TPM hierarchies is already owned"),
			),
			sb_preinstall.ErrInsufficientDMAProtection,
		},
	}

	errorDetails, err := secboot.UnwrapPreinstallCheckError(compoundError)
	c.Assert(err, ErrorMatches, `cannot unwrap error of unexpected type \*errors\.errorString \(the platform firmware indicates that DMA protections are insufficient\)`)
	c.Assert(errorDetails, IsNil)
}

func (s *preinstallSuite) TestUnwrapPreinstallCheckErrorFailCompoundWrapsNil(c *C) {
	compoundError := &CompoundPreinstallCheckError{nil}

	errorDetails, err := secboot.UnwrapPreinstallCheckError(compoundError)
	c.Assert(err, ErrorMatches, "compound error does not wrap any error")
	c.Assert(errorDetails, IsNil)
}

func (s *preinstallSuite) TestUnwrapPreinstallCheckErrorSingle(c *C) {
	singleError := sb_preinstall.NewWithKindAndActionsError(
		sb_preinstall.ErrorKindTPMDeviceDisabled,
		nil,
		[]sb_preinstall.Action{sb_preinstall.ActionRebootToFWSettings},
		errors.New("error with TPM2 device: TPM2 device is present but is currently disabled by the platform firmware"),
	)

	errorDetails, err := secboot.UnwrapPreinstallCheckError(singleError)
	c.Assert(err, IsNil)
	c.Assert(errorDetails, DeepEquals, []secboot.PreinstallErrorDetails{
		{
			Kind:    "tpm-device-disabled",
			Message: "error with TPM2 device: TPM2 device is present but is currently disabled by the platform firmware",
			Actions: []string{"reboot-to-fw-settings"},
		},
	})

	jsn, err := json.MarshalIndent(errorDetails, "", "  ")
	c.Assert(err, IsNil)
	const expectedJson = `[
  {
    "kind": "tpm-device-disabled",
    "message": "error with TPM2 device: TPM2 device is present but is currently disabled by the platform firmware",
    "actions": [
      "reboot-to-fw-settings"
    ]
  }
]`
	c.Assert(string(jsn), Equals, expectedJson)
}

func (s *preinstallSuite) TestUnwrapPreinstallCheckErrorFailSingleUnexpectedType(c *C) {
	errorDetails, err := secboot.UnwrapPreinstallCheckError(sb_preinstall.ErrInsufficientDMAProtection)
	c.Assert(err, ErrorMatches, `cannot unwrap error of unexpected type \*errors\.errorString \(the platform firmware indicates that DMA protections are insufficient\)`)
	c.Assert(errorDetails, IsNil)
}

func (s *preinstallSuite) testPreinstallCheckConfig(c *C, isVM, permitVM bool) {
	cmdExit := `exit 1`
	if isVM {
		cmdExit = `exit 0`
	}
	systemdCmd := testutil.MockCommand(c, "systemd-detect-virt", cmdExit)
	defer systemdCmd.Restore()

	restore := secboot.MockSbPreinstallNewRunChecksContext(
		func(initialFlags sb_preinstall.CheckFlags, loadedImages []sb_efi.Image, profileOpts sb_preinstall.PCRProfileOptionsFlags) *sb_preinstall.RunChecksContext {
			if permitVM {
				c.Assert(initialFlags, Equals, sb_preinstall.PermitVirtualMachine|sb_preinstall.PermitVARSuppliedDrivers)
			} else {
				c.Assert(initialFlags, Equals, sb_preinstall.PermitVARSuppliedDrivers)
			}
			c.Assert(profileOpts, Equals, sb_preinstall.PCRProfileOptionsDefault)
			c.Assert(loadedImages, IsNil)

			return nil
		})
	defer restore()

	restore = secboot.MockSbPreinstallRun(
		func(checkCtx *sb_preinstall.RunChecksContext, ctx context.Context, action sb_preinstall.Action, args ...any) (*sb_preinstall.CheckResult, error) {
			c.Assert(checkCtx, IsNil)
			c.Assert(ctx, NotNil)
			c.Assert(action, Equals, sb_preinstall.ActionNone)
			c.Assert(args, IsNil)

			return &sb_preinstall.CheckResult{}, nil
		})
	defer restore()

	errorDetails, err := secboot.PreinstallCheck(context.Background(), nil)
	c.Assert(err, IsNil)
	c.Assert(errorDetails, IsNil)
}

func (s *preinstallSuite) TestPreinstallCheckConfig(c *C) {
	testCases := []struct {
		isVM     bool
		permitVM bool
	}{
		{false, false}, // default config
		{true, true},   // modify default config to permit VM
	}

	for _, tc := range testCases {
		s.testPreinstallCheckConfig(c, tc.isVM, tc.permitVM)
	}
}

func (s *preinstallSuite) testPreinstallCheck(c *C, detectErrors, failUnwrap bool) {
	bootImagePaths := []string{
		"/cdrom/EFI/boot/bootXXX.efi",
		"/cdrom/EFI/boot/grubXXX.efi",
		"/cdrom/casper/vmlinuz",
	}

	systemdCmd := testutil.MockCommand(c, "systemd-detect-virt", "exit 1")
	defer systemdCmd.Restore()

	restore := secboot.MockSbPreinstallNewRunChecksContext(
		func(initialFlags sb_preinstall.CheckFlags, loadedImages []sb_efi.Image, profileOpts sb_preinstall.PCRProfileOptionsFlags) *sb_preinstall.RunChecksContext {
			c.Assert(initialFlags, Equals, sb_preinstall.PermitVARSuppliedDrivers)
			c.Assert(profileOpts, Equals, sb_preinstall.PCRProfileOptionsDefault)
			c.Assert(loadedImages, HasLen, len(bootImagePaths))
			for i, image := range loadedImages {
				c.Check(image.String(), Equals, bootImagePaths[i])
			}

			return &sb_preinstall.RunChecksContext{}
		})
	defer restore()

	restore = secboot.MockSbPreinstallRun(
		func(checkCtx *sb_preinstall.RunChecksContext, ctx context.Context, action sb_preinstall.Action, args ...any) (*sb_preinstall.CheckResult, error) {
			c.Assert(checkCtx, NotNil)
			c.Assert(checkCtx.Errors(), IsNil)
			c.Assert(checkCtx.Result(), IsNil)
			c.Assert(ctx, NotNil)
			c.Assert(action, Equals, sb_preinstall.ActionNone)
			c.Assert(args, IsNil)

			if detectErrors {
				return nil, &CompoundPreinstallCheckError{
					[]error{
						sb_preinstall.NewWithKindAndActionsError(
							sb_preinstall.ErrorKindTPMHierarchiesOwned,
							sb_preinstall.TPM2OwnedHierarchiesError{WithAuthValue: tpm2.HandleList{tpm2.HandleLockout},
								WithAuthPolicy: tpm2.HandleList{tpm2.HandleOwner}},
							[]sb_preinstall.Action{sb_preinstall.ActionRebootToFWSettings},
							errors.New("error with TPM2 device: one or more of the TPM hierarchies is already owned"),
						),
						sb_preinstall.NewWithKindAndActionsError(
							sb_preinstall.ErrorKindTPMDeviceLockout,
							&sb_preinstall.TPMDeviceLockoutArgs{IntervalDuration: 7200000000000, TotalDuration: 230400000000000},
							[]sb_preinstall.Action{sb_preinstall.ActionRebootToFWSettings},
							errors.New("error with TPM2 device: TPM is in DA lockout mode"),
						),
					},
				}
			} else if failUnwrap {
				return nil, sb_preinstall.ErrInsufficientDMAProtection
			} else {
				return &sb_preinstall.CheckResult{
					Warnings: &CompoundPreinstallCheckError{
						[]error{
							errors.New("warning 1"),
							errors.New("warning 2"),
						},
					},
				}, nil
			}
		})
	defer restore()

	logbuf, restore := logger.MockLogger()
	defer restore()

	errorDetails, err := secboot.PreinstallCheck(context.Background(), bootImagePaths)
	if detectErrors {
		c.Assert(err, IsNil)
		c.Assert(logbuf.String(), Equals, "")
		c.Assert(errorDetails, DeepEquals, []secboot.PreinstallErrorDetails{
			{
				Kind:    "tpm-hierarchies-owned",
				Message: "error with TPM2 device: one or more of the TPM hierarchies is already owned",
				Args: map[string]json.RawMessage{
					"with-auth-value":  json.RawMessage(`[1073741834]`),
					"with-auth-policy": json.RawMessage(`[1073741825]`),
				},
				Actions: []string{"reboot-to-fw-settings"},
			},
			{
				Kind:    "tpm-device-lockout",
				Message: "error with TPM2 device: TPM is in DA lockout mode",
				Args: map[string]json.RawMessage{
					"interval-duration": json.RawMessage(`7200000000000`),
					"total-duration":    json.RawMessage(`230400000000000`),
				},
				Actions: []string{"reboot-to-fw-settings"},
			},
		})
	} else if failUnwrap {
		c.Assert(err, ErrorMatches, `cannot unwrap error of unexpected type \*errors\.errorString \(the platform firmware indicates that DMA protections are insufficient\)`)
		c.Assert(errorDetails, IsNil)
	} else {
		c.Assert(err, IsNil)
		c.Assert(errorDetails, IsNil)
		c.Assert(logbuf.String(), testutil.Contains, "preinstall check warning: warning 1")
		c.Assert(logbuf.String(), testutil.Contains, "preinstall check warning: warning 2")
	}
}

func (s *preinstallSuite) TestPreinstallCheckWithWarningsAndErrors(c *C) {
	detectErrors := false // warnings and no errors
	s.testPreinstallCheck(c, detectErrors, false)
	detectErrors = true // errors and no warnings
	s.testPreinstallCheck(c, detectErrors, false)
}

func (s *preinstallSuite) TestPreinstallCheckFailUnwrap(c *C) {
	failUnwrap := true
	s.testPreinstallCheck(c, false, failUnwrap)
}
