-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathroot.go
More file actions
169 lines (155 loc) · 6.6 KB
/
root.go
File metadata and controls
169 lines (155 loc) · 6.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
// Package cmd implements the CobraCLI commands for the methodaws CLI. Subcommands for the CLI should all live within
// this package. Logic should be delegated to internal packages and functions to keep the CLI commands clean and
// focused on CLI I/O.
package cmd
import (
"errors"
"strings"
"time"
"github.com/Method-Security/methodaws/internal/config"
"github.com/Method-Security/methodaws/utils"
"github.com/Method-Security/pkg/signal"
"github.com/Method-Security/pkg/writer"
"github.com/aws/aws-sdk-go-v2/aws"
awsconfig "github.com/aws/aws-sdk-go-v2/config"
"github.com/palantir/pkg/datetime"
"github.com/palantir/witchcraft-go-logging/wlog/svclog/svc1log"
// Import wlog-zap for its side effects, initializing the zap logger
_ "github.com/palantir/witchcraft-go-logging/wlog-zap"
"github.com/spf13/cobra"
)
// MethodAws is the main struct that holds the root command and all subcommands that are used throughout execution
// of the CLI. It is also responsible for holding the AWS configuration, Output configuration, and Output signal
// for use by subcommands. The output signal is used to write the output of the command to the desired output format
// after the execution of the invoked commands Run function.
type MethodAws struct {
Version string
RootFlags config.RootFlags
OutputConfig writer.OutputConfig
OutputSignal signal.Signal
AwsConfig *aws.Config
RootCmd *cobra.Command
}
// NewMethodAws returns a new MethodAws struct with the provided version string. The MethodAws struct is used to
// initialize the root command and all subcommands that are used throughout execution of the CLI.
// We pass the version command in here from the main.go file, where we set the version string during the build process.
func NewMethodAws(version string) *MethodAws {
startedAt := datetime.DateTime(time.Now())
methodAws := MethodAws{
Version: version,
RootFlags: config.RootFlags{
Quiet: false,
Verbose: false,
Regions: []string{},
},
OutputConfig: writer.NewOutputConfig(nil, writer.NewFormat(writer.SIGNAL)),
OutputSignal: signal.NewSignal(nil, &startedAt, nil, 0, nil),
AwsConfig: nil,
}
return &methodAws
}
// Helper function to set up common configurations
func (a *MethodAws) setupCommonConfig(cmd *cobra.Command, outputFormat string, outputFile string, authed bool) error {
var err error
// Set up output configuration FIRST - before anything that might fail
// This ensures we get properly formatted error output even if AWS auth fails
format, err := validateOutputFormat(outputFormat)
if err != nil {
return err
}
var outputFilePointer *string
if outputFile != "" {
outputFilePointer = &outputFile
} else {
outputFilePointer = nil
}
a.OutputConfig = writer.NewOutputConfig(outputFilePointer, format)
cmd.SetContext(svc1log.WithLogger(cmd.Context(), config.InitializeLogging(cmd, &a.RootFlags)))
if authed {
awsConfig, err := awsconfig.LoadDefaultConfig(cmd.Context())
if err != nil {
return err
}
a.AwsConfig = &awsConfig
if len(a.RootFlags.Regions) != 0 {
a.RootFlags.Regions, err = utils.GetAWSRegions(cmd.Context(), *a.AwsConfig, a.RootFlags.Regions)
if err != nil {
a.OutputSignal.Status = 1
a.OutputSignal.ErrorMessage = aws.String("No valid AWS regions specified")
return errors.New("no valid AWS regions specified")
}
} else {
a.RootFlags.Regions = utils.GetGeneralRegionsList()
}
} else {
a.RootFlags.Regions = utils.GetRegionsToCheck(cmd.Context(), a.RootFlags.Regions)
}
return nil
}
// InitRootCommand initializes the root command for the methodaws CLI. This command is used to set the global flags
// that are used by all subcommands, such as the region, output format, and output file. It also initializes the
// version command that prints the version of the CLI.
// Critically, this sets the PersistentPreRunE and PersistentPostRunE functions that are inherited by most subcommands.
// The PersistentPreRunE function is used to validate the region flag and set the AWS configuration. The PersistentPostRunE
// function is used to write the output of the command to the desired output format after the execution of the invoked
// command's Run function.
func (a *MethodAws) InitRootCommand() {
var outputFormat string
var outputFile string
a.RootCmd = &cobra.Command{
Use: "methodaws",
Short: "Audit AWS resources",
Long: "Audit AWS resources",
SilenceUsage: true,
PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
// Use the variables directly - Cobra automatically updates them when flags are parsed
return a.setupCommonConfig(cmd, outputFormat, outputFile, true)
},
PersistentPostRunE: func(cmd *cobra.Command, _ []string) error {
completedAt := datetime.DateTime(time.Now())
a.OutputSignal.CompletedAt = &completedAt
return writer.Write(
a.OutputSignal.Content,
a.OutputConfig,
&a.OutputSignal.StartedAt,
a.OutputSignal.CompletedAt,
a.OutputSignal.Status,
a.OutputSignal.ErrorMessage,
)
},
}
a.RootCmd.PersistentFlags().BoolVarP(&a.RootFlags.Quiet, "quiet", "q", false, "Suppress output")
a.RootCmd.PersistentFlags().BoolVarP(&a.RootFlags.Verbose, "verbose", "v", false, "Verbose output")
a.RootCmd.PersistentFlags().StringArrayVarP(&a.RootFlags.Regions, "regions", "r", []string{}, "AWS Regions to search for resources. You can specify multiple regions by providing the flag multiple times. If blank, will search all regions.")
a.RootCmd.PersistentFlags().StringVarP(&outputFile, "output-file", "f", "", "Path to output file. If blank, will output to STDOUT")
a.RootCmd.PersistentFlags().StringVarP(&outputFormat, "output", "o", "signal", "Output format (signal, json, yaml). Default value is signal")
versionCmd := &cobra.Command{
Use: "version",
Short: "Print the version number of methodaws",
PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
return nil
},
Run: func(cmd *cobra.Command, args []string) {
cmd.Println(a.Version)
},
PersistentPostRunE: func(cmd *cobra.Command, _ []string) error {
return nil
},
}
a.RootCmd.AddCommand(versionCmd)
}
// A utility function to validate that the provided output format is one of the supported formats: json, yaml, signal.
func validateOutputFormat(output string) (writer.Format, error) {
var format writer.FormatValue
switch strings.ToLower(output) {
case "json":
format = writer.JSON
case "yaml":
return writer.Format{}, errors.New("yaml output format is not supported for methodaws")
case "signal":
format = writer.SIGNAL
default:
return writer.Format{}, errors.New("invalid output format. Valid formats are: json, yaml, signal")
}
return writer.NewFormat(format), nil
}