Adds `save_pe32` functionality to UTK so I can more easily experiment with UEFI firmware from the command line without legacy UEFItool

savepe32.go

// Copyright 2018 the LinuxBoot Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package visitors

import (
	"errors"
	"os"

	"github.com/linuxboot/fiano/pkg/uefi"
)

// SavePE32 saves the first PE32 matching the query
type SavePE32 struct {
	// Input
	Predicate func(f uefi.Firmware) bool
	OutPE32   string

	// Output
	Matches []uefi.Firmware
}

// Run wraps Visit and performs some setup and teardown tasks.
func (v *SavePE32) Run(f uefi.Firmware) error {
	// Run "find" to generate a list of matches to replace.
	find := Find{
		Predicate: v.Predicate,
	}
	if err := find.Run(f); err != nil {
		return err
	}

	// Use this list of matches
	v.Matches = find.Matches
	if len(find.Matches) == 0 {
		return errors.New("no matches found for replacement")
	}
	if len(find.Matches) > 1 {
		return errors.New("multiple matches found! There can be only one. Use find to list all matches")
	}

	for _, m := range v.Matches {
		if err := m.Apply(v); err != nil {
			return err
		}
	}
	return nil
}

// Visit applies the Extract visitor to any Firmware type.
func (v *SavePE32) Visit(f uefi.Firmware) error {
	switch f := f.(type) {
	case *uefi.File:
		return f.ApplyChildren(v)
	case *uefi.Section:
		if f.Header.Type == uefi.SectionTypePE32 {
			return os.WriteFile( v.OutPE32, f.Buf()[4:], 0666 )
		}
		return nil
	default:
		// Must be applied to a File to have any effect.
		return nil
	}
}

func init() {
	RegisterCLI("save_pe32", "saves a pe32 given a GUID and output path", 2, func(args []string) (uefi.Visitor, error) {
		pred, err := FindFilePredicate(args[0])
		if err != nil {
			return nil, err
		}

		filename := args[1]

		// Find all the matching files and save them
		return &SavePE32{
			Predicate: pred,
			OutPE32:   filename,
		}, nil
	})
}