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 }