带地理位置信息的traceroute

traceroute_with_geo.go

package main

import (
	"errors"
	"fmt"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"net"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
)

var (
	ErrRemoteAddr = errors.New("RemoteAddr error")
)

type TraceRoute struct {
	LocalAddr  string //default 0.0.0.0
	RemoteAddr string
	MaxTTL     int           //default 30
	Timeout    time.Duration //default 3 sec
}
type Result struct {
	ID  int
	IP  string
	RTT time.Duration
}

func (r Result) String() string {
	if r.IP == "*" {
		return fmt.Sprintf("%d\t%v", r.ID, r.IP)
	}
	return fmt.Sprintf("%d\t%v\t%.2fms", r.ID, r.IP, (r.RTT.Seconds() * 1000.0))
}

func New(remote string) *TraceRoute {
	return &TraceRoute{
		LocalAddr:  "0.0.0.0",
		RemoteAddr: remote,
		MaxTTL:     30,
		Timeout:    3 * time.Second,
	}
}

func (t *TraceRoute) Do(callback func(Result)) ([]Result, error) {
	ips, err := net.LookupIP(t.RemoteAddr)
	if err != nil {
		return nil, err
	}
	var dst net.IPAddr
	for _, ip := range ips {
		if ip.To4() != nil {
			dst.IP = ip
			break
		}
	}
	if dst.IP == nil {
		return nil, ErrRemoteAddr
	}
	c, err := net.ListenPacket("ip4:icmp", t.LocalAddr)
	if err != nil {
		return nil, err
	}
	defer c.Close()
	p := ipv4.NewPacketConn(c)
	if err := p.SetControlMessage(ipv4.FlagTTL|ipv4.FlagSrc|ipv4.FlagDst|ipv4.FlagInterface, true); err != nil {
		return nil, err
	}
	wm := icmp.Message{
		Type: ipv4.ICMPTypeEcho, Code: 0, Body: &icmp.Echo{ID: os.Getpid() & 0xffff, Data: []byte("R-U-OK?")},
	}
	rb := make([]byte, 1500)
	var result []Result
	for i := 1; i < t.MaxTTL; i++ {
		wm.Body.(*icmp.Echo).Seq = i
		wb, err := wm.Marshal(nil)
		if err != nil {
			return result, err
		}
		if err := p.SetTTL(i); err != nil {
			return result, err
		}
		begin := time.Now()
		if _, err := p.WriteTo(wb, nil, &dst); err != nil {
			return result, err
		}
		if err := p.SetReadDeadline(time.Now().Add(t.Timeout)); err != nil {
			return result, err
		}
		n, _, peer, err := p.ReadFrom(rb)
		if err != nil {
			if err, ok := err.(net.Error); ok && err.Timeout() {
				result = append(result, Result{ID: i, IP: "*"})
				continue
			}
			return result, err
		}
		rm, err := icmp.ParseMessage(1, rb[:n])
		if err != nil {
			return result, err
		}
		rtt := time.Since(begin)
		switch rm.Type {
		case ipv4.ICMPTypeTimeExceeded:
			if callback != nil {
				callback(Result{ID: i, IP: peer.String(), RTT: rtt})
			}
			result = append(result, Result{ID: i, IP: peer.String(), RTT: rtt})
		case ipv4.ICMPTypeEchoReply:
			if callback != nil {
				callback(Result{ID: i, IP: peer.String(), RTT: rtt})
			}
			result = append(result, Result{ID: i, IP: peer.String(), RTT: rtt})
			return result, nil
		default:
			if callback != nil {
				callback(Result{ID: i, IP: "*"})
			}
			result = append(result, Result{ID: i, IP: "*"})
		}
	}
	return result, nil
}

type GeoIpInfo struct {
	Ip            string
	CountryCode   string
	CountryCode3  string
	Country       string
	ContinentCode string
	Region        string
	RegionCode    string
	City          string
	PostalCode    string
	DmaCode       string
	AreaCode      string
	Organization  string
	Timezone      string
	Offset        string
	Latitude      float32
	Longitude     float32
}

func IpGeoAddress(addr string) string {
	//采用WEB版的数据，API版本的数据不全
	doc, err := goquery.NewDocument("https://ip.sb/ip/" + addr)
	if err != nil {
		return err.Error()
	}

	segments := []string{}

	r := regexp.MustCompile("[ \r\n\t]{2,}")

	if segment := doc.Find("td.proto_isp").Text(); segment != "" {
		segment = r.ReplaceAllString(strings.Trim(segment, " \t\r\n"), "")
		if segment != "None" && segment != "Unkown" {
			segments = append(segments, segment)
		}
	}
	if segment := doc.Find("td.proto_hostname span").Text(); segment != "" {
		segment = r.ReplaceAllString(strings.Trim(segment, " \t\r\n"), "")
		if segment != "None" && segment != "Unkown" {
			segments = append(segments, segment)
		}
	}
	if segment := doc.Find("td.proto_location").Text(); segment != "" {
		segment = r.ReplaceAllString(strings.Trim(segment, " \t\r\n"), "")
		if segment != "None" && segment != "Unkown" {
			segments = append(segments, segment)
		}
	}

	if len(segments) == 0 {
		return "Unkown"
	}
	return strings.Join(segments, " ")
}

func Callback(result Result) {
	fmt.Printf("%2d %-15s %s\n", result.ID, result.IP, IpGeoAddress(result.IP))
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: " + os.Args[0] + " IP/Domain")
		return
	}

	t := New(os.Args[1])
	_, err := t.Do(Callback)
	if err != nil {
		fmt.Println(err)
		return
	}
}