From ef328a6a481f5ec69bf5f409b6cd0b2061477a2b Mon Sep 17 00:00:00 2001
From: Evan Chen <evan.yf.chen@ibm.com>
Date: Fri, 15 Oct 2021 18:57:15 +0800
Subject: [PATCH] update

---
 Makefile |   2 +
 go.sum   |   0
 main.go  | 147 +++++++++++++++++++++++++++++++++++++++++++------------
 3 files changed, 118 insertions(+), 31 deletions(-)
 create mode 100644 go.sum

diff --git a/Makefile b/Makefile
index 98f0c2c..3ab8044 100644
--- a/Makefile
+++ b/Makefile
@@ -1,2 +1,4 @@
 build:
 	go build -o bin/myip
+
+# release:
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..e69de29
diff --git a/main.go b/main.go
index 5085ab2..52aa5d5 100644
--- a/main.go
+++ b/main.go
@@ -4,10 +4,12 @@ import (
 	"flag"
 	"fmt"
 	"io/ioutil"
+	"log"
 	"net"
 	"net/http"
 	"os"
 	"strings"
+	"text/tabwriter"
 )
 
 var (
@@ -17,16 +19,30 @@ var (
 	envP       string
 	host       string
 	envHost    string
+	showVer    bool
+	showHeader bool
+	showAll    bool
+	showTab    bool
+	// iface string
+	// envIface string
 )
 
+var Version = "v0.1.0"
+
 func init() {
 	envP = os.Getenv("MYIP_PORT")
 	envHost = os.Getenv("MYIP_HOST")
 
+	flag.BoolVar(&showVer, "v", false, "show version")
 	flag.BoolVar(&serverMode, "s", false, "run in server mode")
+	flag.BoolVar(&showAll, "a", false, "show all interface")
+	flag.BoolVar(&showHeader, "t", false, "show header")
+	flag.BoolVar(&showTab, "T", false, "show tab")
 	flag.StringVar(&host, "u", "kumoly.io", "hostname to ask for")
 	flag.StringVar(&port, "p", "5080", "port to listen|connect")
 	flag.StringVar(&addr, "b", "0.0.0.0", "the address the server binds on")
+
+	// flag.StringVar(&port, "p", "5080", "port to listen|connect")
 }
 
 func usage() {
@@ -36,6 +52,10 @@ func usage() {
 
 func main() {
 	flag.Parse()
+	if showVer {
+		fmt.Println(Version)
+		return
+	}
 	if envP != "" {
 		port = envP
 	}
@@ -44,39 +64,104 @@ func main() {
 	}
 
 	if serverMode {
-		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-			ip := r.Header.Get("X-Real-Ip")
-			if ip == "" {
-				ip = r.Header.Get("X-Forwarded-For")
-			}
-			if ip == "" {
-				var err error
-				ip, _, err = net.SplitHostPort(r.RemoteAddr)
-				if err != nil {
-					ip = r.RemoteAddr
-				}
-			}
-			fmt.Fprint(w, ip+"\n")
-		})
-
-		err := http.ListenAndServe(addr+":"+port, nil)
-		if err != nil {
-			fmt.Println(err)
-			os.Exit(1)
-		}
-
+		StartServer()
 	} else {
-		res, err := http.Get("http://" + host + ":" + port)
-		if err != nil {
-			fmt.Println(err)
-			os.Exit(2)
+		var wFlag uint = 0
+		if showTab {
+			wFlag = tabwriter.Debug
 		}
-		body, err := ioutil.ReadAll(res.Body)
-		res.Body.Close()
-		if err != nil {
-			fmt.Println(err)
-			os.Exit(3)
+		w := tabwriter.NewWriter(os.Stdout, 0, 0, 1, ' ', wFlag)
+		if showHeader {
+			fmt.Fprintln(w, "\tInterface\tIP\t")
 		}
-		fmt.Println(strings.TrimRight(string(body), "\r\n"))
+		fmt.Fprintln(w, fmt.Sprintf("\t%v\t%v\t", "PublicIP", GetPublicIP()))
+		if showAll {
+			for k, v := range GetLocalIP() {
+				fmt.Fprintln(w, fmt.Sprintf("\t%v\t%v\t", strings.Trim(k, " "), v))
+			}
+		}
+		w.Flush()
 	}
 }
+
+func StartServer() {
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		ip := r.Header.Get("X-Real-Ip")
+		if ip == "" {
+			ip = r.Header.Get("X-Forwarded-For")
+		}
+		if ip == "" {
+			var err error
+			ip, _, err = net.SplitHostPort(r.RemoteAddr)
+			if err != nil {
+				ip = r.RemoteAddr
+			}
+		}
+		fmt.Fprint(w, ip+"\n")
+	})
+
+	http.HandleFunc("/version", func(rw http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(rw, Version)
+	})
+
+	err := http.ListenAndServe(addr+":"+port, nil)
+	if err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
+}
+
+func GetLocalIP() map[string]string {
+	ret := map[string]string{}
+	ifaces, err := net.Interfaces()
+	if err != nil {
+		log.Println(err)
+		return ret
+	}
+
+	for _, i := range ifaces {
+		addrs, err := i.Addrs()
+		if err != nil {
+			log.Println(err)
+			return ret
+		}
+
+		for _, addr := range addrs {
+			var ip net.IP
+			switch v := addr.(type) {
+			case *net.IPNet:
+				ip = v.IP
+				if ip.IsLoopback() {
+					continue
+				}
+				if ip.To4() == nil {
+					continue
+				}
+			case *net.IPAddr:
+				continue
+				// ip = v.IP
+				// if ip.IsLoopback() {
+				// 	continue
+				// }
+			}
+			ret[i.Name] = ip.String()
+			break
+		}
+	}
+	return ret
+}
+
+func GetPublicIP() string {
+	res, err := http.Get("http://" + host + ":" + port)
+	if err != nil {
+		log.Println(err)
+		return ""
+	}
+	body, err := ioutil.ReadAll(res.Body)
+	res.Body.Close()
+	if err != nil {
+		log.Println(err)
+		return ""
+	}
+	return strings.TrimRight(string(body), "\r\n")
+}