github.com/canonical/ubuntu-image@v0.0.0-20240430122802-2202fe98b290/cmd/ubuntu-image/main.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"os"
     7  
     8  	"github.com/jessevdk/go-flags"
     9  
    10  	"github.com/canonical/ubuntu-image/internal/commands"
    11  	"github.com/canonical/ubuntu-image/internal/helper"
    12  	"github.com/canonical/ubuntu-image/internal/statemachine"
    13  )
    14  
    15  // Version holds the ubuntu-image version number
    16  // this is usually overridden at build time
    17  var Version string = ""
    18  
    19  // helper variables for unit testing
    20  var osExit = os.Exit
    21  var captureStd = helper.CaptureStd
    22  
    23  var stateMachineLongDesc = `Options for controlling the internal state machine.
    24  Other than -w, these options are mutually exclusive. When -u or -t is given,
    25  the state machine can be resumed later with -r, but -w must be given in that
    26  case since the state is saved in a ubuntu-image.json file in the working directory.`
    27  
    28  func initStateMachine(imageType string, commonOpts *commands.CommonOpts, stateMachineOpts *commands.StateMachineOpts, ubuntuImageCommand *commands.UbuntuImageCommand) (statemachine.SmInterface, error) {
    29  	var stateMachine statemachine.SmInterface
    30  	switch imageType {
    31  	case "snap":
    32  		stateMachine = &statemachine.SnapStateMachine{
    33  			Opts: ubuntuImageCommand.Snap.SnapOptsPassed,
    34  			Args: ubuntuImageCommand.Snap.SnapArgsPassed,
    35  		}
    36  	case "classic":
    37  		stateMachine = &statemachine.ClassicStateMachine{
    38  			Args: ubuntuImageCommand.Classic.ClassicArgsPassed,
    39  		}
    40  	case "pack":
    41  		stateMachine = &statemachine.PackStateMachine{
    42  			Opts: ubuntuImageCommand.Pack.PackOptsPassed,
    43  		}
    44  	default:
    45  		return nil, fmt.Errorf("unsupported command\n")
    46  	}
    47  
    48  	stateMachine.SetCommonOpts(commonOpts, stateMachineOpts)
    49  
    50  	return stateMachine, nil
    51  }
    52  
    53  func executeStateMachine(sm statemachine.SmInterface) error {
    54  	if err := sm.Setup(); err != nil {
    55  		return err
    56  	}
    57  
    58  	if err := sm.Run(); err != nil {
    59  		return err
    60  	}
    61  
    62  	if err := sm.Teardown(); err != nil {
    63  		return err
    64  	}
    65  
    66  	return nil
    67  }
    68  
    69  // unhidePackOpts make pack options visible in help if the pack command is used
    70  // This should be removed when the pack command is made visible to everyone
    71  func unhidePackOpts(parser *flags.Parser) {
    72  	// Save given options before removing them temporarily
    73  	// otherwise the help will be displayed twice
    74  	opts := parser.Options
    75  	parser.Options = 0
    76  	defer func() { parser.Options = opts }()
    77  	// parse once to determine the active command
    78  	// we do not care about error here since we will reparse again
    79  	_, _ = parser.Parse() // nolint: errcheck
    80  
    81  	if parser.Active != nil {
    82  		if parser.Active.Name == "pack" {
    83  			parser.Active.Hidden = false
    84  		}
    85  	}
    86  }
    87  
    88  // parseFlags parses received flags and returns error code accordingly
    89  func parseFlags(parser *flags.Parser, restoreStdout, restoreStderr func(), stdout, stderr io.Reader, resume, version bool) (error, int) {
    90  	if _, err := parser.Parse(); err != nil {
    91  		if e, ok := err.(*flags.Error); ok {
    92  			switch e.Type {
    93  			case flags.ErrHelp:
    94  				restoreStdout()
    95  				restoreStderr()
    96  				readStdout, err := io.ReadAll(stdout)
    97  				if err != nil {
    98  					fmt.Printf("Error reading from stdout: %s\n", err.Error())
    99  					return err, 1
   100  				}
   101  				fmt.Println(string(readStdout))
   102  				return e, 0
   103  			case flags.ErrCommandRequired:
   104  				// if --resume was given, this is not an error
   105  				if !resume && !version {
   106  					restoreStdout()
   107  					restoreStderr()
   108  					readStderr, err := io.ReadAll(stderr)
   109  					if err != nil {
   110  						fmt.Printf("Error reading from stderr: %s\n", err.Error())
   111  						return err, 1
   112  					}
   113  					fmt.Printf("Error: %s\n", string(readStderr))
   114  					return e, 1
   115  				}
   116  			default:
   117  				restoreStdout()
   118  				restoreStderr()
   119  				fmt.Printf("Error: %s\n", err.Error())
   120  				return e, 1
   121  			}
   122  		}
   123  	}
   124  	return nil, 0
   125  }
   126  
   127  func main() { //nolint: gocyclo
   128  	commonOpts := new(commands.CommonOpts)
   129  	stateMachineOpts := new(commands.StateMachineOpts)
   130  	ubuntuImageCommand := new(commands.UbuntuImageCommand)
   131  
   132  	// set up the go-flags parser for command line options
   133  	parser := flags.NewParser(ubuntuImageCommand, flags.Default)
   134  	_, err := parser.AddGroup("State Machine Options", stateMachineLongDesc, stateMachineOpts)
   135  	if err != nil {
   136  		fmt.Printf("Error: %s\n", err.Error())
   137  		osExit(1)
   138  		return
   139  	}
   140  	_, err = parser.AddGroup("Common Options", "Options common to every command", commonOpts)
   141  	if err != nil {
   142  		fmt.Printf("Error: %s\n", err.Error())
   143  		osExit(1)
   144  		return
   145  	}
   146  
   147  	// go-flags can be overzealous about printing errors that aren't actually errors
   148  	// so we capture stdout/stderr while parsing and later decide whether to print
   149  	stdout, restoreStdout, err := captureStd(&os.Stdout)
   150  	if err != nil {
   151  		fmt.Printf("Failed to capture stdout: %s\n", err.Error())
   152  		osExit(1)
   153  		return
   154  	}
   155  	defer restoreStdout()
   156  
   157  	stderr, restoreStderr, err := captureStd(&os.Stderr)
   158  	if err != nil {
   159  		fmt.Printf("Failed to capture stderr: %s\n", err.Error())
   160  		osExit(1)
   161  		return
   162  	}
   163  	defer restoreStderr()
   164  
   165  	unhidePackOpts(parser)
   166  
   167  	// Parse the options provided and handle specific errors
   168  	err, code := parseFlags(parser, restoreStdout, restoreStderr, stdout, stderr, stateMachineOpts.Resume, commonOpts.Version)
   169  	if err != nil {
   170  		osExit(code)
   171  		return
   172  	}
   173  
   174  	// restore stdout
   175  	restoreStdout()
   176  	restoreStderr()
   177  
   178  	// in case user only requested version number, print and exit
   179  	if commonOpts.Version {
   180  		// we expect Version to be supplied at build time or fetched from the snap environment
   181  		if Version == "" {
   182  			Version = os.Getenv("SNAP_VERSION")
   183  		}
   184  		fmt.Printf("ubuntu-image %s\n", Version)
   185  		osExit(0)
   186  		return
   187  	}
   188  
   189  	var imageType string
   190  	if parser.Command.Active != nil {
   191  		imageType = parser.Command.Active.Name
   192  	}
   193  
   194  	// init the state machine
   195  	sm, err := initStateMachine(imageType, commonOpts, stateMachineOpts, ubuntuImageCommand)
   196  	if err != nil {
   197  		fmt.Printf("Error: %s\n", err.Error())
   198  		osExit(1)
   199  		return
   200  	}
   201  
   202  	// let the state machine handle the image build
   203  	err = executeStateMachine(sm)
   204  	if err != nil {
   205  		fmt.Printf("Error: %s\n", err.Error())
   206  		osExit(1)
   207  		return
   208  	}
   209  }