// This file is part of nullboot
// Copyright 2021 Canonical Ltd.
// SPDX-License-Identifier: GPL-3.0-only

// Package efibootmgr contains a boot management library
package efibootmgr

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"path"

	"github.com/canonical/go-efilib"
	efi_linux "github.com/canonical/go-efilib/linux"
)

const (
	maxBootEntries = 65535 // Maximum number of boot entries we can hold
	bootCurrentStr = "BootCurrent"
	bootNextStr    = "BootNext"
	bootOrderStr   = "BootOrder"
)
const invalidBootNumber = -1
const defaultBootVariableAttributes = efi.AttributeNonVolatile | efi.AttributeBootserviceAccess | efi.AttributeRuntimeAccess

type BootVarNotFoundError struct {
	Message string
	Err     error
}

func (e BootVarNotFoundError) Error() string {
	return e.Message
}

func (e BootVarNotFoundError) Unwrap() error {
	return e.Err
}

// BootEntryVariable defines a boot entry variable
type BootEntryVariable struct {
	BootNumber int                    // number of the Boot variable, for example, for Boot0004 this is 4
	Data       []byte                 // the data of the variable
	Attributes efi.VariableAttributes // any attributes set on the variable
	LoadOption *efi.LoadOption        // the data of the variable parsed as a load option, if it is a valid load option
}

// NewBootEntryVariable creates a BootEntryVariable derived from a BootEntry,
// given a boot number and device path.
func NewBootEntryVariable(entry BootEntry, bootNum int, dp efi.DevicePath) (BootEntryVariable, error) {
	loadOption := newEFILoadOption(entry, dp)
	data, err := loadOption.Bytes()
	if err != nil {
		return BootEntryVariable{}, fmt.Errorf("cannot encode boot entry data: %w", err)
	}

	entryVar := BootEntryVariable{
		BootNumber: bootNum,
		Data:       data,
		Attributes: defaultBootVariableAttributes,
		LoadOption: loadOption,
	}

	return entryVar, nil
}

// BootManager manages the boot device selection menu entries (Boot0000...BootFFFF).
type BootManager struct {
	efivars        EFIVariables              // EFIVariables implementation
	entries        map[int]BootEntryVariable // The Boot<number> variables
	bootOrder      []int                     // The BootOrder variable, parsed
	bootCurrent    int                       // The BootCurrent variable, parsed
	bootNext       int                       // The BootNext variable, parsed
	bootOrderAttrs efi.VariableAttributes    // The attributes of BootOrder variable
}

// NewBootManagerFromSystem returns a new BootManager object, initialized with the system state.
func NewBootManagerFromSystem() (BootManager, error) {
	return NewBootManagerForVariables(RealEFIVariables{})
}

// NewBootManagerForVariables returns a boot manager for the given EFIVariables manager
func NewBootManagerForVariables(efivars EFIVariables) (BootManager, error) {
	var err error
	bm := BootManager{}
	bm.efivars = efivars
	bm.bootNext = invalidBootNumber

	if !VariablesSupported(efivars) {
		return BootManager{}, fmt.Errorf("Variables not supported")
	}

	if bootCurrentBytes, _, err := bm.efivars.GetVariable(efi.GlobalVariable, bootCurrentStr); err != nil {
		switch bm.efivars.(type) {
		// BootCurrent is a volatile entry that is created upon a
		// successful boot into some boot entry. Since MockEFIVariables is
		// not representing a live host, BootCurrent does not need to exist
		case *MockEFIVariables:
			log.Printf("Could not read %s variable, populating with default (%d), error was: %v\n", bootCurrentStr, invalidBootNumber, err)
			bm.bootCurrent = invalidBootNumber
		default:
			return BootManager{}, fmt.Errorf("could not read %s variable, error was: %w\n", bootCurrentStr, err)
		}
	} else {
		bm.bootCurrent = int(binary.LittleEndian.Uint16(bootCurrentBytes[0:2]))
	}

	// It's possible that BootNext is active from user input or a previous iteration of nullboot
	if bootNextBytes, _, err := bm.efivars.GetVariable(efi.GlobalVariable, bootNextStr); err != nil {
		if len(bootNextBytes) != 0 {
			bm.bootNext = int(binary.LittleEndian.Uint16(bootNextBytes[0:2]))
		}
	}

	bootOrderBytes, bootOrderAttrs, err := bm.efivars.GetVariable(efi.GlobalVariable, bootOrderStr)
	if err != nil {
		log.Printf("Could not read %s variable, populating with default, error was: %v", bootOrderStr, err)
		bootOrderBytes = nil
		bootOrderAttrs = efi.AttributeNonVolatile | efi.AttributeBootserviceAccess | efi.AttributeRuntimeAccess
	}
	bm.bootOrder = make([]int, len(bootOrderBytes)/2)
	bm.bootOrderAttrs = bootOrderAttrs
	for i := 0; i < len(bootOrderBytes); i += 2 {
		// FIXME: It's probably not valid to assume little-endian here?
		bm.bootOrder[i/2] = int(binary.LittleEndian.Uint16(bootOrderBytes[i : i+2]))
	}

	bm.entries = make(map[int]BootEntryVariable)
	names, err := GetVariableNames(bm.efivars, efi.GlobalVariable)
	if err != nil {
		return BootManager{}, fmt.Errorf("cannot obtain list of global variables: %v", err)
	}
	for _, name := range names {
		var entry BootEntryVariable
		if parsed, err := fmt.Sscanf(name, "Boot%04X", &entry.BootNumber); len(name) != 8 || parsed != 1 || err != nil {
			continue
		}
		entry.Data, entry.Attributes, err = bm.efivars.GetVariable(efi.GlobalVariable, name)
		if err != nil {
			return BootManager{}, fmt.Errorf("cannot read %s: %v", name, err)
		}
		entry.LoadOption, err = efi.ReadLoadOption(bytes.NewReader(entry.Data))
		if err != nil {
			log.Printf("Invalid boot entry Boot%04X: %s\n", entry.BootNumber, err)
		}

		bm.entries[entry.BootNumber] = entry
	}

	return bm, nil
}

// NextFreeEntry returns the number of the next free Boot variable.
func (bm *BootManager) NextFreeEntry() (int, error) {
	for i := 0; i < maxBootEntries; i++ {
		if _, ok := bm.entries[i]; !ok {
			return i, nil
		}
	}

	return invalidBootNumber, fmt.Errorf("Maximum number of boot entries exceeded")
}

// FindOrCreateEntry finds a matching entry in the boot device selection menu,
// or creates one if it is missing.
//
// It returns the number of the entry created, or invalidBootNumber on failure, with error set.
//
// The argument relativeTo specifies the directory entry.Filename is in.
func (bm *BootManager) FindOrCreateEntry(entry BootEntry, relativeTo string) (int, error) {
	bootEntryVar, err := bm.FindBootEntryVariable(entry, relativeTo)
	if err == nil {
		return bootEntryVar.BootNumber, nil
	} else {
		var notFoundError *BootVarNotFoundError
		if !errors.As(err, &notFoundError) {
			err = fmt.Errorf("unable to determine if boot entry %s exists: %w", entry.Label, err)
			return invalidBootNumber, err
		}
	}

	// Entry was not found; generate requisite data
	dpStr := path.Join(relativeTo, entry.Filename)
	dp, err := bm.efivars.NewFileDevicePath(dpStr, efi_linux.ShortFormPathHD)
	if err != nil {
		return invalidBootNumber, fmt.Errorf("unable to derive device path %s: %w", dpStr, err)
	}
	freeEntryNumber, err := bm.NextFreeEntry()
	if err != nil {
		return invalidBootNumber, fmt.Errorf("unable to generate a new boot number: %w", err)
	}

	entryVar, err := NewBootEntryVariable(entry, freeEntryNumber, dp)
	if err != nil {
		return invalidBootNumber, fmt.Errorf("unable to create the boot entry variable: %w", err)
	}

	// Entry needs to be added as a boot entry
	err = bm.RegisterBootEntryVariable(entryVar)
	if err != nil {
		return entryVar.BootNumber, fmt.Errorf(
			"unable to set environment variables for %s: %w",
			toEFIBootEntryFormat(entryVar.BootNumber),
			err,
		)
	}
	bm.entries[freeEntryNumber] = entryVar
	return entryVar.BootNumber, nil
}

// FindBootEntryVariable finds a matching BootEntryVariable from the boot
// device selection menu associated to the input BootEntry in a relative
// directory.
//
// The argument relativeTo specifies the directory entry.Filename is in.
//
// If a BootEntryVariable cannot be found, returns BootVarNotFoundError.
func (bm *BootManager) FindBootEntryVariable(entry BootEntry, relativeTo string) (BootEntryVariable, error) {
	dp, err := bm.efivars.NewFileDevicePath(path.Join(relativeTo, entry.Filename), efi_linux.ShortFormPathHD)
	if err != nil {
		return BootEntryVariable{}, fmt.Errorf("unable to derive device path: %w", err)
	}

	loadOption := newEFILoadOption(entry, dp)
	data, err := loadOption.Bytes()
	if err != nil {
		return BootEntryVariable{}, fmt.Errorf("unable to encode load option for %s: %w", entry.Label, err)
	}

	for _, existingVar := range bm.entries {
		if bytes.Equal(existingVar.Data, data) && existingVar.Attributes == defaultBootVariableAttributes {
			return existingVar, nil
		}
	}

	err = fmt.Errorf("unable to find %s in the registered EFI boot variables: %w", entry.Label, err)
	notFoundErr := BootVarNotFoundError{
		Message: err.Error(),
		Err:     err,
	}
	return BootEntryVariable{}, &notFoundErr
}

// RegisterBootEntryVariable registers a system boot variable to the
// bm.efivars and bm.entries.
//
// It returns an error should there be an issue registering the system variable.
func (bm *BootManager) RegisterBootEntryVariable(entryVar BootEntryVariable) error {
	variable := toEFIBootEntryFormat(entryVar.BootNumber)

	if err := bm.efivars.SetVariable(efi.GlobalVariable, variable, entryVar.Data, entryVar.Attributes); err != nil {
		return err
	}
	bm.entries[entryVar.BootNumber] = entryVar
	return nil
}

// DeleteEntry deletes an entry and updates the cached boot order.
//
// The boot order still needs to be committed afterwards. It is not written back immediately,
// as there will usually be multiple places to update boot order, and we can coalesce those
// writes. We still have to update the boot order though, such that when we delete an entry
// and then create a new one with the same number we don't accidentally have the new one in
// the order.
func (bm *BootManager) DeleteEntry(bootNum int) error {
	variable := toEFIBootEntryFormat(bootNum)
	if _, ok := bm.entries[bootNum]; !ok {
		return fmt.Errorf("Tried deleting a non-existing variable %s", variable)
	}

	if err := DelVariable(bm.efivars, efi.GlobalVariable, variable); err != nil {
		return err
	}
	delete(bm.entries, bootNum)

	// Deleting BootCurrent is not a good idea. It can put your system at
	// risk. apt won't delete the user's BootCurrent artifacts, so this has
	// to be an action caused by a bug or a mistaken user; however, if the
	// user truly deleted all artifacts associated to BootCurrent, the FDE
	// key is likely resealed at this point without those artifacts, and
	// removing this entry won't lead to any **more** mistakes.
	if bootNum == bm.bootCurrent {
		log.Printf("the entry associated to BootCurrent (%s) has been deleted\n", toEFIBootEntryFormat(bootNum))
	} else if bootNum == bm.bootNext {
		log.Printf("the entry associated to BootNext (%s) has been deleted and un-set\n", toEFIBootEntryFormat(bootNum))
		if err := DelVariable(bm.efivars, efi.GlobalVariable, bootNextStr); err != nil {
			return err
		}

		bm.bootNext = invalidBootNumber
	}

	var newOrder []int

	for _, orderEntry := range bm.bootOrder {
		if orderEntry != bootNum {
			newOrder = append(newOrder, orderEntry)
		}

	}

	bm.bootOrder = newOrder

	return nil
}

// PrependAndSetBootOrder commits a new boot order or returns an error.
//
// The boot order specified is prepended to the existing one, and the order
// is deduplicated before committing.
func (bm *BootManager) PrependAndSetBootOrder(head []int) error {
	var newOrder []int

	// Combine head with existing boot order, filter out duplicates and non-existing entries
	for _, num := range append(append([]int(nil), head...), bm.bootOrder...) {
		isDuplicate := false
		for _, otherNum := range newOrder {
			if otherNum == num {
				isDuplicate = true
			}
		}
		if _, ok := bm.entries[num]; ok && !isDuplicate {
			newOrder = append(newOrder, num)
		}
	}

	// Encode the boot order to bytes
	var output []byte
	for _, num := range newOrder {
		numBytes := toEFIBootEntryBytes(num)
		output = append(output, numBytes[0], numBytes[1])
	}

	// Set the boot order and update our cache
	if err := bm.efivars.SetVariable(efi.GlobalVariable, bootOrderStr, output, bm.bootOrderAttrs); err != nil {
		return err
	}

	bm.bootOrder = newOrder
	return nil

}

// toEFIBootEntryFormat generates the Boot Name associated to the specified integer
// and formats it to match the EFI Boot Entry style:
// i.e. toEFIBootEntryFormat(1) -> "Boot0001", toEFIBootEntryFormat(43) -> "Boot002B", ...
func toEFIBootEntryFormat(bootNum int) string {
	return fmt.Sprintf("Boot%04X", bootNum)
}

// SetBootNext sets the system BootNext variable to the encoded value of
// the input boot number.
//
// Returns an error in the event that the variable cannot be set.
func (bm *BootManager) SetBootNext(bootNum int) error {
	bootEntryName := toEFIBootEntryFormat(bootNum)
	if _, _, err := bm.efivars.GetVariable(efi.GlobalVariable, bootEntryName); err != nil {
		return fmt.Errorf("unable to find %s: %w", bootEntryName, err)
	}
	bootNumBytes := toEFIBootEntryBytes(bootNum)
	if err := bm.efivars.SetVariable(efi.GlobalVariable, bootNextStr, bootNumBytes, bm.bootOrderAttrs); err != nil {
		return err
	}

	bm.bootNext = bootNum
	return nil
}

// newEFILoadOption derives a standardized LoadOption from a specified entry
// and device path.
//
// It returns a pointer to the new LoadOption.
func newEFILoadOption(entry BootEntry, dp efi.DevicePath) *efi.LoadOption {
	optionalData := new(bytes.Buffer)
	binary.Write(optionalData, binary.LittleEndian, efi.ConvertUTF8ToUCS2(entry.Options+"\x00"))
	loadOption := &efi.LoadOption{
		Attributes:   efi.LoadOptionActive,
		Description:  entry.Label,
		FilePath:     dp,
		OptionalData: optionalData.Bytes()}

	return loadOption
}

// toEFIBootEntryBytes converts a boot number integer into the system format.
//
// Returns a 2-length []byte of little-endian oriented uint16s.
func toEFIBootEntryBytes(bootNum int) []byte {
	numBytes := make([]byte, 2)
	binary.LittleEndian.PutUint16(numBytes[:], uint16(bootNum))
	return numBytes
}
