github.com/linuxboot/fiano@v1.2.0/pkg/visitors/count.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  	"encoding/json"
     9  	"fmt"
    10  	"io"
    11  	"os"
    12  	"strings"
    13  
    14  	"github.com/linuxboot/fiano/pkg/uefi"
    15  )
    16  
    17  // Count counts the number of each firmware type.
    18  type Count struct {
    19  	// Optionally write result as JSON.
    20  	W io.Writer `json:"-"`
    21  
    22  	// Output
    23  	FirmwareTypeCount map[string]int
    24  	FileTypeCount     map[string]int
    25  	SectionTypeCount  map[string]int
    26  }
    27  
    28  // Run wraps Visit and performs some setup and teardown tasks.
    29  func (v *Count) Run(f uefi.Firmware) error {
    30  	v.FirmwareTypeCount = map[string]int{}
    31  	v.FileTypeCount = map[string]int{}
    32  	v.SectionTypeCount = map[string]int{}
    33  
    34  	if err := f.Apply(v); err != nil {
    35  		return err
    36  	}
    37  
    38  	if v.W != nil {
    39  		b, err := json.MarshalIndent(v, "", "\t")
    40  		if err != nil {
    41  			return err
    42  		}
    43  		_, err = fmt.Fprintln(v.W, string(b))
    44  		return err
    45  	}
    46  	return nil
    47  }
    48  
    49  // Visit applies the Count visitor to any Firmware type.
    50  func (v *Count) Visit(f uefi.Firmware) error {
    51  	incr := func(m *map[string]int, key string) {
    52  		if n, ok := (*m)[key]; ok {
    53  			(*m)[key] = n + 1
    54  		} else {
    55  			(*m)[key] = 1
    56  		}
    57  	}
    58  
    59  	incr(&v.FirmwareTypeCount, strings.TrimPrefix(fmt.Sprintf("%T", f), "*uefi."))
    60  	if file, ok := f.(*uefi.File); ok {
    61  		incr(&v.FileTypeCount, file.Type)
    62  	}
    63  	if sec, ok := f.(*uefi.Section); ok {
    64  		incr(&v.SectionTypeCount, sec.Type)
    65  	}
    66  	return f.ApplyChildren(v)
    67  }
    68  
    69  func init() {
    70  	RegisterCLI("count", "count the number of each firmware type", 0, func(args []string) (uefi.Visitor, error) {
    71  		return &Count{
    72  			W: os.Stdout,
    73  		}, nil
    74  	})
    75  }