diff --git a/cmd/analysis.go b/cmd/analysis.go index cab46c2..f2f41e3 100644 --- a/cmd/analysis.go +++ b/cmd/analysis.go @@ -14,6 +14,7 @@ package cmd import ( + "github.com/VirusTotal/vt-cli/utils" "regexp" "github.com/spf13/cobra" @@ -49,7 +50,10 @@ func NewAnalysisCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("analyses/%s", args, re) + return p.GetAndPrintObjects( + "analyses/%s", + utils.StringReaderFromCmdArgs(args), + re) }, } diff --git a/cmd/domain.go b/cmd/domain.go index 8ed5c22..4b2665d 100644 --- a/cmd/domain.go +++ b/cmd/domain.go @@ -14,6 +14,7 @@ package cmd import ( + "github.com/VirusTotal/vt-cli/utils" "github.com/spf13/cobra" ) @@ -46,8 +47,10 @@ func NewDomainCmd() *cobra.Command { if err != nil { return err } - p.GetAndPrintObjects("domains/%s", args, nil) - return nil + return p.GetAndPrintObjects( + "domains/%s", + utils.StringReaderFromCmdArgs(args), + nil) }, } diff --git a/cmd/file.go b/cmd/file.go index ae04266..bbc6e81 100644 --- a/cmd/file.go +++ b/cmd/file.go @@ -14,6 +14,7 @@ package cmd import ( + "github.com/VirusTotal/vt-cli/utils" "regexp" "github.com/spf13/cobra" @@ -49,7 +50,10 @@ func NewFileCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("files/%s", args, re) + return p.GetAndPrintObjects( + "files/%s", + utils.StringReaderFromCmdArgs(args), + re) }, } diff --git a/cmd/group.go b/cmd/group.go index c57119b..f6e8a4e 100644 --- a/cmd/group.go +++ b/cmd/group.go @@ -14,6 +14,7 @@ package cmd import ( + "github.com/VirusTotal/vt-cli/utils" "github.com/spf13/cobra" ) @@ -35,7 +36,10 @@ func NewGroupCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("groups/%s", args, nil) + return p.GetAndPrintObjects( + "groups/%s", + utils.StringReaderFromCmdArgs(args), + nil) }, } diff --git a/cmd/hunting.go b/cmd/hunting.go index c9f2f68..5392885 100644 --- a/cmd/hunting.go +++ b/cmd/hunting.go @@ -17,6 +17,7 @@ import ( "bufio" "errors" "fmt" + "github.com/VirusTotal/vt-cli/utils" "os" "regexp" "strconv" @@ -124,7 +125,10 @@ func NewHuntingNotificationCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("intelligence/hunting_notifications/%s", args, re) + return p.GetAndPrintObjects( + "intelligence/hunting_notifications/%s", + utils.StringReaderFromCmdArgs(args), + re) }, } @@ -352,7 +356,10 @@ func NewHuntingRulesetCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("intelligence/hunting_rulesets/%s?relationships=owner,editors", args, re) + return p.GetAndPrintObjects( + "intelligence/hunting_rulesets/%s?relationships=owner,editors", + utils.StringReaderFromCmdArgs(args), + re) }, } diff --git a/cmd/ip.go b/cmd/ip.go index daff2c2..e28a4d5 100644 --- a/cmd/ip.go +++ b/cmd/ip.go @@ -14,6 +14,7 @@ package cmd import ( + "github.com/VirusTotal/vt-cli/utils" "regexp" "github.com/spf13/cobra" @@ -47,7 +48,10 @@ func NewIPCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("ip_addresses/%s", args, re) + return p.GetAndPrintObjects( + "ip_addresses/%s", + utils.StringReaderFromCmdArgs(args), + re) }, } diff --git a/cmd/monitor.go b/cmd/monitor.go index 3b93c71..8498867 100644 --- a/cmd/monitor.go +++ b/cmd/monitor.go @@ -458,7 +458,9 @@ func NewMonitorCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("monitor/items/%s", args, re) + return p.GetAndPrintObjects("monitor/items/%s", + utils.StringReaderFromCmdArgs(args), + re) }, } diff --git a/cmd/retrohunt.go b/cmd/retrohunt.go index a15d71a..d1eae80 100644 --- a/cmd/retrohunt.go +++ b/cmd/retrohunt.go @@ -15,6 +15,7 @@ package cmd import ( "fmt" + "github.com/VirusTotal/vt-cli/utils" "io/ioutil" "os" "regexp" @@ -324,7 +325,10 @@ func NewRetrohuntCmd() *cobra.Command { if err != nil { return err } - return p.GetAndPrintObjects("intelligence/retrohunt_jobs/%s", args, re) + return p.GetAndPrintObjects( + "intelligence/retrohunt_jobs/%s", + utils.StringReaderFromCmdArgs(args), + re) }, } diff --git a/cmd/url.go b/cmd/url.go index 6b76a8f..cdafe01 100644 --- a/cmd/url.go +++ b/cmd/url.go @@ -15,7 +15,7 @@ package cmd import ( "encoding/base64" - + "github.com/VirusTotal/vt-cli/utils" "github.com/spf13/cobra" ) @@ -48,11 +48,12 @@ func NewURLCmd() *cobra.Command { if err != nil { return err } - ids := make([]string, len(args)) - for i, arg := range args { - ids[i] = base64.RawURLEncoding.EncodeToString([]byte(arg)) - } - return p.GetAndPrintObjects("urls/%s", ids, nil) + r := utils.NewMappedStringReader( + utils.StringReaderFromCmdArgs(args), + func (url string) string { + return base64.RawURLEncoding.EncodeToString([]byte(url)) + }) + return p.GetAndPrintObjects("urls/%s", r, nil) }, } diff --git a/cmd/user.go b/cmd/user.go index 55329b9..880e378 100644 --- a/cmd/user.go +++ b/cmd/user.go @@ -16,6 +16,7 @@ package cmd import ( "strings" + "github.com/VirusTotal/vt-cli/utils" "github.com/spf13/cobra" ) @@ -45,7 +46,9 @@ func NewUserCmd() *cobra.Command { "api_quota_group", "intelligence_quota_group", "monitor_quota_group", - }, ","), args, nil) + }, ","), + utils.StringReaderFromCmdArgs(args), + nil) }, } diff --git a/go.mod b/go.mod index 2828e0a..e39a001 100644 --- a/go.mod +++ b/go.mod @@ -15,7 +15,6 @@ require ( github.com/mattn/go-isatty v0.0.10 // indirect github.com/mattn/go-runewidth v0.0.5 // indirect github.com/mitchellh/go-homedir v1.1.0 - github.com/mitchellh/gox v1.0.1 // indirect github.com/pelletier/go-toml v1.6.0 // indirect github.com/plusvic/go-ansi v0.0.0-20180516115420-9879244c4340 github.com/spf13/afero v1.2.2 // indirect diff --git a/go.sum b/go.sum index cdf8f9f..65048cc 100644 --- a/go.sum +++ b/go.sum @@ -56,8 +56,6 @@ github.com/gosuri/uitable v0.0.3/go.mod h1:tKR86bXuXPZazfOTG1FIzvjIdXzd0mo4Vtn16 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs= github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk= github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY= -github.com/hashicorp/go-version v1.0.0 h1:21MVWPKDphxa7ineQQTrCU5brh7OuVVAzGOCnnCPtE8= -github.com/hashicorp/go-version v1.0.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA= github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM= @@ -88,10 +86,6 @@ github.com/mattn/go-runewidth v0.0.5/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= -github.com/mitchellh/gox v1.0.1 h1:x0jD3dcHk9a9xPSDN6YEL4xL6Qz0dvNYm8yZqui5chI= -github.com/mitchellh/gox v1.0.1/go.mod h1:ED6BioOGXMswlXa2zxfh/xdd5QhwYliBFn9V18Ap4z4= -github.com/mitchellh/iochan v1.0.0 h1:C+X3KsSTLFVBr/tK1eYN/vs4rJcvsiLU338UhYPJWeY= -github.com/mitchellh/iochan v1.0.0/go.mod h1:JwYml1nuB7xOzsp52dPpHFffvOCDupsG0QubkSMEySY= github.com/mitchellh/mapstructure v1.1.2 h1:fmNYVwqnSfB9mZU6OS2O6GsXM+wcskZDuKQzvN1EDeE= github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y= github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= diff --git a/utils/printer.go b/utils/printer.go index 8c27ba2..47b6002 100644 --- a/utils/printer.go +++ b/utils/printer.go @@ -124,16 +124,7 @@ func (p *Printer) PrintObject(obj *vt.Object) error { // items from the args slice. If args contains a single "-" string, the args are // read from stdin one per line. If argRe is non-nil, only args that match the // regular expression are used and the rest are discarded. -func (p *Printer) GetAndPrintObjects(endpoint string, args []string, argRe *regexp.Regexp) error { - - var r StringReader - - if len(args) == 1 && args[0] == "-" { - r = NewStringIOReader(os.Stdin) - } else { - r = NewStringArrayReader(args) - } - +func (p *Printer) GetAndPrintObjects(endpoint string, r StringReader, argRe *regexp.Regexp) error { if argRe != nil { r = NewFilteredStringReader(r, argRe) } diff --git a/utils/string_reader.go b/utils/string_reader.go index c82b225..7bd33bf 100644 --- a/utils/string_reader.go +++ b/utils/string_reader.go @@ -16,6 +16,7 @@ package utils import ( "bufio" "io" + "os" "regexp" "strings" ) @@ -86,9 +87,9 @@ func NewFilteredStringReader(r StringReader, re *regexp.Regexp) *FilteredStringR return &FilteredStringReader{r: r, re: re} } -// ReadString reads strings from the the underlying StringReader and returns -// the first one that matches the regular expression specified while creating -// the FilteredStringReader. If no more strings can be read err is io.EOF. +// ReadString reads strings from the underlying StringReader and returns the +// first one that matches the regular expression specified while creating the +// FilteredStringReader. If no more strings can be read err is io.EOF. func (f *FilteredStringReader) ReadString() (s string, err error) { for s, err = f.r.ReadString(); s != "" || err == nil; s, err = f.r.ReadString() { if f.re.MatchString(s) { @@ -97,3 +98,38 @@ func (f *FilteredStringReader) ReadString() (s string, err error) { } return s, err } + + +// MappedStringReader reads strings from a StringReader and call a map function +// that transforms the strings in some other string. +type MappedStringReader struct { + r StringReader + mapFn func(string) string +} + +// NewMappedStringReader creates a new MappedStringReader that reads strings from +// r and can call mapFn for transforming the string before returning it. +func NewMappedStringReader(r StringReader, mapFn func(string) string) *MappedStringReader { + return &MappedStringReader{r:r, mapFn: mapFn} +} + +// ReadString reads strings from the underlying StringReader and can call the +// map function associated to the MappedStringReader with that string, then +// returns the result produced by the map function. +func (m *MappedStringReader) ReadString() (s string, err error) { + if s, err = m.r.ReadString(); err == nil { + return m.mapFn(s), nil + } + return s, err +} + +// StringReaderFromCmdArgs returns a string reader for reading the arguments +// passed in the command line. If the arguments consists in single hypen "-", +// they are read from stdin. +func StringReaderFromCmdArgs(args []string) StringReader { + if len(args) == 1 && args[0] == "-" { + return NewStringIOReader(os.Stdin) + } else { + return NewStringArrayReader(args) + } +}