1717package defaultsubcommand
1818
1919import (
20+ "bytes"
2021 "fmt"
2122 "io"
2223 "os"
24+ "path/filepath"
25+ "regexp"
2326
24- nvctkConfig "github.com/NVIDIA/nvidia-container-toolkit/internal/config"
27+ "github.com/NVIDIA/nvidia-container-toolkit/internal/config"
2528 "github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
2629 "github.com/urfave/cli/v2"
2730)
@@ -32,7 +35,9 @@ type command struct {
3235
3336// options stores the subcommand options
3437type options struct {
35- output string
38+ config string
39+ output string
40+ inPlace bool
3641}
3742
3843// NewCommand constructs a default command with the specified logger
@@ -61,9 +66,20 @@ func (m command) build() *cli.Command {
6166 }
6267
6368 c .Flags = []cli.Flag {
69+ & cli.StringFlag {
70+ Name : "config" ,
71+ Usage : "Specify the config file to process; The contents of this file overrides the default config" ,
72+ Destination : & opts .config ,
73+ },
74+ & cli.BoolFlag {
75+ Name : "in-place" ,
76+ Aliases : []string {"i" },
77+ Usage : "Modify the config file in-place" ,
78+ Destination : & opts .inPlace ,
79+ },
6480 & cli.StringFlag {
6581 Name : "output" ,
66- Usage : "Specify the file to output the generated configuration for to. If this is '' the configuration is ouput to STDOUT. " ,
82+ Usage : "Specify the output file to write to; If not specified, the output is written to stdout " ,
6783 Destination : & opts .output ,
6884 },
6985 }
@@ -72,31 +88,96 @@ func (m command) build() *cli.Command {
7288}
7389
7490func (m command ) validateFlags (c * cli.Context , opts * options ) error {
91+ if opts .inPlace {
92+ if opts .output != "" {
93+ return fmt .Errorf ("cannot specify both --in-place and --output" )
94+ }
95+ opts .output = opts .config
96+ }
7597 return nil
7698}
7799
78100func (m command ) run (c * cli.Context , opts * options ) error {
79- defaultConfig , err := nvctkConfig .GetDefaultConfigToml ()
101+ if err := opts .ensureOutputFolder (); err != nil {
102+ return fmt .Errorf ("unable to create output directory: %v" , err )
103+ }
104+
105+ contents , err := opts .getFormattedConfig ()
80106 if err != nil {
81- return fmt .Errorf ("unable to get default config : %v" , err )
107+ return fmt .Errorf ("unable to fix comments : %v" , err )
82108 }
83109
110+ if _ , err := opts .Write (contents ); err != nil {
111+ return fmt .Errorf ("unable to write to output: %v" , err )
112+ }
113+
114+ return nil
115+ }
116+
117+ // getFormattedConfig returns the default config formatted as required from the specified config file.
118+ // The config is then formatted as required.
119+ // No indentation is used and comments are modified so that there is no space
120+ // after the '#' character.
121+ func (opts options ) getFormattedConfig () ([]byte , error ) {
122+ cfg , err := config .Load (opts .config )
123+ if err != nil {
124+ return nil , fmt .Errorf ("unable to load or create config: %v" , err )
125+ }
126+
127+ buffer := bytes .NewBuffer (nil )
128+
129+ if _ , err := cfg .Save (buffer ); err != nil {
130+ return nil , fmt .Errorf ("unable to save config: %v" , err )
131+ }
132+ return fixComments (buffer .Bytes ())
133+ }
134+
135+ func fixComments (contents []byte ) ([]byte , error ) {
136+ r , err := regexp .Compile (`(\n*)\s*?#\s*(\S.*)` )
137+ if err != nil {
138+ return nil , fmt .Errorf ("unable to compile regexp: %v" , err )
139+ }
140+ replaced := r .ReplaceAll (contents , []byte ("$1#$2" ))
141+
142+ return replaced , nil
143+ }
144+
145+ func (opts options ) outputExists () (bool , error ) {
146+ if opts .output == "" {
147+ return false , nil
148+ }
149+ _ , err := os .Stat (opts .output )
150+ if err == nil {
151+ return true , nil
152+ } else if ! os .IsNotExist (err ) {
153+ return false , fmt .Errorf ("unable to stat output file: %v" , err )
154+ }
155+ return false , nil
156+ }
157+
158+ func (opts options ) ensureOutputFolder () error {
159+ if opts .output == "" {
160+ return nil
161+ }
162+ if dir := filepath .Dir (opts .output ); dir != "" {
163+ return os .MkdirAll (dir , 0755 )
164+ }
165+ return nil
166+ }
167+
168+ // Write writes the contents to the output file specified in the options.
169+ func (opts options ) Write (contents []byte ) (int , error ) {
84170 var output io.Writer
85171 if opts .output == "" {
86172 output = os .Stdout
87173 } else {
88174 outputFile , err := os .Create (opts .output )
89175 if err != nil {
90- return fmt .Errorf ("unable to create output file: %v" , err )
176+ return 0 , fmt .Errorf ("unable to create output file: %v" , err )
91177 }
92178 defer outputFile .Close ()
93179 output = outputFile
94180 }
95181
96- _ , err = defaultConfig .WriteTo (output )
97- if err != nil {
98- return fmt .Errorf ("unable to write to output: %v" , err )
99- }
100-
101- return nil
182+ return output .Write (contents )
102183}
0 commit comments