add short flags

go.mod

@@ -1,7 +1,10 @@
-module main
+module github.com/byReqz/probehost2
 
 go 1.17
 
 require github.com/sirupsen/logrus v1.8.1
 
-require golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e // indirect
+require (
+	github.com/spf13/pflag v1.0.5 // indirect
+	golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e // indirect
+)

go.sum

@@ -4,6 +4,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
 github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4=

main.go

@@ -6,20 +6,23 @@ import (
   "strings"
   "net/http"
   "net"
-  "flag"
 
   log "github.com/sirupsen/logrus"
+  flag "github.com/spf13/pflag"
+
 )
 
 var logstdout = log.New()
 var logfile = log.New()
 
 var listenport int
+var usexforwardedfor bool
 
 func init() {
   var logfilepath string
-  flag.StringVar(&logfilepath, "logfilepath", "probehost2.log", "sets the output file for the log")
-  flag.IntVar(&listenport, "port", 8000, "sets the port to listen on")
+  flag.StringVarP(&logfilepath, "logfilepath", "o","probehost2.log", "sets the output file for the log")
+  flag.IntVarP(&listenport, "port", "p", 8000, "sets the port to listen on")
+  flag.BoolVarP(&usexforwardedfor, "use-x-forwarded-for", "x", true, "specifies whether to show x-forwarded-for or the requesting IP")
   flag.Parse()
 
   logstdout.SetFormatter(&log.TextFormatter{
@@ -82,24 +85,38 @@ func parseopts(options []string, cmdopts map[string]string) []string{
   return opts
 }
 
-func ping(w http.ResponseWriter, req *http.Request) {
+func prerunner(req *http.Request, cmd string, cmdopts map[string]string, defaultopts []string) string{
   geturl := strings.Split(req.URL.String(), "/")
   targets := strings.Split(geturl[2], ",")
   hosts := validatehosts(targets)
-  cmdopts := map[string]string{"4": "-4", "6": "-6", "d": "-D", "n": "-n", "v": "-v", "c1": "-c1", "c5": "-c5", "c10": "-c10"}
   var opts []string
-  opts = append(opts, "-c10") // add default options
-  if len(geturl) > 3 { 
+  opts = append(opts, defaultopts...)
+  if len(geturl) > 3 && len(geturl[3]) > 0 { 
     options := strings.Split(geturl[3], ",")
     opts = append(opts, parseopts(options, cmdopts)...)
   }
   var res string
   var args []string
+  var remoteaddr string
+  if req.Header.Get("X-Forwarded-For") != "" && usexforwardedfor != false {
+    remoteaddr = req.Header.Get("X-Forwarded-For")
+  } else {
+    remoteaddr = req.RemoteAddr
+  }
   for _, host := range hosts {
     args = append(args, opts...)
     args = append(args, host)
-    res = fmt.Sprint(res, runner(req.RemoteAddr, "ping", args...), "\n")
+    res = fmt.Sprint(res, runner(remoteaddr, cmd, args...), "\n")
   }
+  return res
+}
+
+func ping(w http.ResponseWriter, req *http.Request) {
+  cmd := "ping"
+  cmdopts := map[string]string{"4": "-4", "6": "-6", "d": "-D", "n": "-n", "v": "-v", "c1": "-c1", "c5": "-c5", "c10": "-c10"}
+  var defaultopts []string
+  defaultopts = append(defaultopts, "-c10")
+  res := prerunner(req, cmd, cmdopts, defaultopts)
   if strings.TrimSpace(res) == "" {
     fmt.Fprintln(w, http.StatusInternalServerError)
   } else {
@@ -108,23 +125,11 @@ func ping(w http.ResponseWriter, req *http.Request) {
 }
 
 func mtr(w http.ResponseWriter, req *http.Request) {
-  geturl := strings.Split(req.URL.String(), "/")
-  targets := strings.Split(geturl[2], ",")
-  hosts := validatehosts(targets)
+  cmd := "mtr"
   cmdopts := map[string]string{"4": "-4", "6": "-6", "u": "-u", "t": "-T", "e": "-e", "x": "-x", "n": "-n", "b": "-b", "z": "-z", "c1": "-c1", "c5": "-c5", "c10": "-c10"}
-  var opts []string
-  opts = append(opts, "-r", "-w", "-c10") // add default options
-  if len(geturl) > 3 { 
-    options := strings.Split(geturl[3], ",")
-    opts = append(opts, parseopts(options, cmdopts)...)
-  }
-  var res string
-  var args []string
-  for _, host := range hosts {
-    args = append(args, opts...)
-    args = append(args, host)
-    res = fmt.Sprint(res, runner(req.RemoteAddr, "mtr", args...), "\n")
-  }
+  var defaultopts []string
+  defaultopts = append(defaultopts, "-r", "-w", "-c10")
+  res := prerunner(req, cmd, cmdopts, defaultopts)
   if strings.TrimSpace(res) == "" {
     fmt.Fprintln(w, http.StatusInternalServerError)
   } else {
@@ -133,23 +138,11 @@ func mtr(w http.ResponseWriter, req *http.Request) {
 }
 
 func traceroute(w http.ResponseWriter, req *http.Request) {
-  geturl := strings.Split(req.URL.String(), "/")
-  targets := strings.Split(geturl[2], ",")
-  hosts := validatehosts(targets)
+  cmd := "traceroute"
   cmdopts := map[string]string{"4": "-4", "6": "-6", "dnf": "-F", "i": "-I", "t": "-T", "n": "-n", "u": "-U", "ul": "-UL", "d": "-D", "b": "--back"}
-  var opts []string
-  //opts = append(opts, "-c10") // no default options for traceroute
-  if len(geturl) > 3 { 
-    options := strings.Split(geturl[3], ",")
-    opts = append(opts, parseopts(options, cmdopts)...)
-  }
-  var res string
-  var args []string
-  for _, host := range hosts {
-    args = append(args, opts...)
-    args = append(args, host)
-    res = fmt.Sprint(res, runner(req.RemoteAddr, "traceroute", args...), "\n")
-  }
+  var defaultopts []string
+  //defaultopts = append(defaultopts) // no default options for traceroute
+  res := prerunner(req, cmd, cmdopts, defaultopts)
   if strings.TrimSpace(res) == "" {
     fmt.Fprintln(w, http.StatusInternalServerError)
   } else {
@@ -161,6 +154,7 @@ func main() {
   http.HandleFunc("/ping/", ping)
   http.HandleFunc("/mtr/", mtr)
   http.HandleFunc("/tracert/", traceroute)
+  http.HandleFunc("/traceroute/", traceroute)
   logstdout.Info("Serving on :", listenport)
   logfile.Info("Serving on :", listenport)
   http.ListenAndServe(fmt.Sprint(":", listenport), nil)