github.com/linuxboot/fiano@v1.2.0/pkg/visitors/replacepe32.go (about)

     1  // Copyright 2018 the LinuxBoot Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package visitors
     6  
     7  import (
     8  	"bytes"
     9  	"errors"
    10  	"os"
    11  
    12  	"github.com/linuxboot/fiano/pkg/uefi"
    13  )
    14  
    15  // ReplacePE32 replaces PE32 sections with NewPE32 for all files matching Predicate.
    16  type ReplacePE32 struct {
    17  	// Input
    18  	Predicate func(f uefi.Firmware) bool
    19  	NewPE32   []byte
    20  
    21  	// Output
    22  	Matches []uefi.Firmware
    23  }
    24  
    25  // Run wraps Visit and performs some setup and teardown tasks.
    26  func (v *ReplacePE32) Run(f uefi.Firmware) error {
    27  	// Check that we're actually replacing with a PE32 image
    28  	if !bytes.HasPrefix(v.NewPE32, []byte("MZ")) {
    29  		return errors.New("supplied binary is not a valid pe32 image")
    30  	}
    31  
    32  	// Run "find" to generate a list of matches to replace.
    33  	find := Find{
    34  		Predicate: v.Predicate,
    35  	}
    36  	if err := find.Run(f); err != nil {
    37  		return err
    38  	}
    39  
    40  	// Use this list of matches for replacing sections.
    41  	v.Matches = find.Matches
    42  	if len(find.Matches) == 0 {
    43  		return errors.New("no matches found for replacement")
    44  	}
    45  	if len(find.Matches) > 1 {
    46  		return errors.New("multiple matches found! There can be only one. Use find to list all matches")
    47  	}
    48  
    49  	for _, m := range v.Matches {
    50  		if err := m.Apply(v); err != nil {
    51  			return err
    52  		}
    53  	}
    54  	return nil
    55  }
    56  
    57  // Visit applies the Extract visitor to any Firmware type.
    58  func (v *ReplacePE32) Visit(f uefi.Firmware) error {
    59  	switch f := f.(type) {
    60  
    61  	case *uefi.File:
    62  		return f.ApplyChildren(v)
    63  
    64  	case *uefi.Section:
    65  		if f.Header.Type == uefi.SectionTypePE32 {
    66  			f.SetBuf(v.NewPE32)
    67  			f.Encapsulated = nil // Should already be empty
    68  			if err := f.GenSecHeader(); err != nil {
    69  				return err
    70  			}
    71  		}
    72  		return f.ApplyChildren(v)
    73  
    74  	default:
    75  		// Must be applied to a File to have any effect.
    76  		return nil
    77  	}
    78  }
    79  
    80  func init() {
    81  	RegisterCLI("replace_pe32", "replace a pe32 given a GUID and new file", 2, func(args []string) (uefi.Visitor, error) {
    82  		pred, err := FindFilePredicate(args[0])
    83  		if err != nil {
    84  			return nil, err
    85  		}
    86  
    87  		filename := args[1]
    88  		newPE32, err := os.ReadFile(filename)
    89  		if err != nil {
    90  			return nil, err
    91  		}
    92  
    93  		// Find all the matching files and replace their inner PE32s.
    94  		return &ReplacePE32{
    95  			Predicate: pred,
    96  			NewPE32:   newPE32,
    97  		}, nil
    98  	})
    99  }