feat(contrib/snmp2cpe): add --port/-P option (future-architect#2046)

budirs86 · Oct 10, 2024 · 2c51bcf · 2c51bcf
1 parent 087b620
commit 2c51bcf
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 5 deletions.
diff --git a/contrib/snmp2cpe/pkg/cmd/v1/v1.go b/contrib/snmp2cpe/pkg/cmd/v1/v1.go
@@ -13,12 +13,14 @@ import (
 
 // SNMPv1Options ...
 type SNMPv1Options struct {
+	Port  uint16
 	Debug bool
 }
 
 // NewCmdV1 ...
 func NewCmdV1() *cobra.Command {
 	opts := &SNMPv1Options{
+		Port:  161,
 		Debug: false,
 	}
 
@@ -28,7 +30,7 @@ func NewCmdV1() *cobra.Command {
 		Example: "$ snmp2cpe v1 192.168.100.1 public",
 		Args:    cobra.ExactArgs(2),
 		RunE: func(_ *cobra.Command, args []string) error {
-			r, err := snmp.Get(gosnmp.Version1, args[0], snmp.WithCommunity(args[1]), snmp.WithDebug(opts.Debug))
+			r, err := snmp.Get(gosnmp.Version1, args[0], snmp.WithCommunity(args[1]), snmp.WithPort(opts.Port), snmp.WithDebug(opts.Debug))
 			if err != nil {
 				return errors.Wrap(err, "failed to snmpget")
 			}
@@ -41,6 +43,7 @@ func NewCmdV1() *cobra.Command {
 		},
 	}
 
+	cmd.Flags().Uint16VarP(&opts.Port, "port", "P", 161, "port")
 	cmd.Flags().BoolVarP(&opts.Debug, "debug", "", false, "debug mode")
 
 	return cmd

diff --git a/contrib/snmp2cpe/pkg/cmd/v2c/v2c.go b/contrib/snmp2cpe/pkg/cmd/v2c/v2c.go
@@ -13,12 +13,14 @@ import (
 
 // SNMPv2cOptions ...
 type SNMPv2cOptions struct {
+	Port  uint16
 	Debug bool
 }
 
 // NewCmdV2c ...
 func NewCmdV2c() *cobra.Command {
 	opts := &SNMPv2cOptions{
+		Port:  161,
 		Debug: false,
 	}
 
@@ -28,7 +30,7 @@ func NewCmdV2c() *cobra.Command {
 		Example: "$ snmp2cpe v2c 192.168.100.1 public",
 		Args:    cobra.ExactArgs(2),
 		RunE: func(_ *cobra.Command, args []string) error {
-			r, err := snmp.Get(gosnmp.Version2c, args[0], snmp.WithCommunity(args[1]), snmp.WithDebug(opts.Debug))
+			r, err := snmp.Get(gosnmp.Version2c, args[0], snmp.WithCommunity(args[1]), snmp.WithPort(opts.Port), snmp.WithDebug(opts.Debug))
 			if err != nil {
 				return errors.Wrap(err, "failed to snmpget")
 			}
@@ -41,6 +43,7 @@ func NewCmdV2c() *cobra.Command {
 		},
 	}
 
+	cmd.Flags().Uint16VarP(&opts.Port, "port", "P", 161, "port")
 	cmd.Flags().BoolVarP(&opts.Debug, "debug", "", false, "debug mode")
 
 	return cmd

diff --git a/contrib/snmp2cpe/pkg/cmd/v3/v3.go b/contrib/snmp2cpe/pkg/cmd/v3/v3.go
@@ -10,12 +10,14 @@ import (
 
 // SNMPv3Options ...
 type SNMPv3Options struct {
+	Port  uint16
 	Debug bool
 }
 
 // NewCmdV3 ...
 func NewCmdV3() *cobra.Command {
 	opts := &SNMPv3Options{
+		Port:  161,
 		Debug: false,
 	}
 
@@ -24,7 +26,7 @@ func NewCmdV3() *cobra.Command {
 		Short:   "snmpget with SNMPv3",
 		Example: "$ snmp2cpe v3",
 		RunE: func(_ *cobra.Command, _ []string) error {
-			_, err := snmp.Get(gosnmp.Version3, "", snmp.WithDebug(opts.Debug))
+			_, err := snmp.Get(gosnmp.Version3, "", snmp.WithPort(opts.Port), snmp.WithDebug(opts.Debug))
 			if err != nil {
 				return errors.Wrap(err, "failed to snmpget")
 			}
@@ -33,6 +35,7 @@ func NewCmdV3() *cobra.Command {
 		},
 	}
 
+	cmd.Flags().Uint16VarP(&opts.Port, "port", "P", 161, "port")
 	cmd.Flags().BoolVarP(&opts.Debug, "debug", "", false, "debug mode")
 
 	return cmd

diff --git a/contrib/snmp2cpe/pkg/snmp/snmp.go b/contrib/snmp2cpe/pkg/snmp/snmp.go
@@ -12,6 +12,7 @@ import (
 
 type options struct {
 	community string
+	port      uint16
 	debug     bool
 }
 
@@ -31,6 +32,17 @@ func WithCommunity(c string) Option {
 	return communityOption(c)
 }
 
+type portOption uint16
+
+func (p portOption) apply(opts *options) {
+	opts.port = uint16(p)
+}
+
+// WithPort ...
+func WithPort(p uint16) Option {
+	return portOption(p)
+}
+
 type debugOption bool
 
 func (d debugOption) apply(opts *options) {
@@ -44,7 +56,11 @@ func WithDebug(d bool) Option {
 
 // Get ...
 func Get(version gosnmp.SnmpVersion, ipaddr string, opts ...Option) (Result, error) {
-	var options options
+	options := options{
+		community: "public",
+		port:      161,
+		debug:     false,
+	}
 	for _, o := range opts {
 		o.apply(&options)
 	}
@@ -53,7 +69,7 @@ func Get(version gosnmp.SnmpVersion, ipaddr string, opts ...Option) (Result, err
 
 	params := &gosnmp.GoSNMP{
 		Target:             ipaddr,
-		Port:               161,
+		Port:               options.port,
 		Version:            version,
 		Timeout:            time.Duration(2) * time.Second,
 		Retries:            3,