package main

import (
	_ "embed"
	"flag"
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
	"time"

	"io/ioutil"
	"net/http"
	"net/url"

	"encoding/json"

	"github.com/aeden/traceroute"
	"github.com/oschwald/maxminddb-golang"
)

var isFinish bool = false
var destIP net.IP

func remove(slice []string, s int) []string {
	if s < len(slice) {
		return append(slice[:s], slice[s+1:]...)
	}
	return slice
}

func removeEmpty(slice []string) (ret []string) {
	for i := 0; i < len(slice); i++ {
		if (len(slice[i])) != 0 {
			ret = append(ret, slice[i])
		}
	}
	return
}

type record struct {
	ASN int    `maxminddb:"autonomous_system_number"`
	ASO string `maxminddb:"autonomous_system_organization"`
}

type IPDetail struct {
	Country  string `json:"country"`
	Province string `json:"province"`
	City     string `json:"city"`
	ISP      string `json:"isp"`
}

type APIResponse struct {
	Code   int      `json:"code"`
	Detail IPDetail `json:"data"`
}

//go:embed `GeoLite2-ASN.mmdb`
var buffer []byte

func parseIPFromMap42(ip net.IP) (record, APIResponse, error) {
	params := url.Values{}

	URL, err := url.Parse("http://ipip.map.dn42/whois")

	if err != nil {
		return record{}, APIResponse{}, err
	}

	params.Set("ip", ip.String())
	params.Set("lang", "cn")

	URL.RawQuery = params.Encode()

	resp, err := http.Get(URL.String())

	if err != nil {
		return record{}, APIResponse{}, err
	}

	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)
	var res struct {
		ASN  string `json:"as"`
		Area string `json:"area"`
	}
	_ = json.Unmarshal(body, &res)

	ASN, err := strconv.Atoi(strings.Replace(res.ASN, "AS", "", -1))
	if err != nil {
		return record{}, APIResponse{}, err
	}

	record := record{
		ASN: ASN,
	}

	APIResponse := APIResponse{
		Code: 0,
		Detail: IPDetail{
			ISP: strings.Join(removeEmpty(remove(remove(strings.Split(res.Area, "\t"), 6), 5)), " "),
		},
	}
	return record, APIResponse, nil
}

func parseIPFromMaxminddb(ip net.IP) record {
	db, err := maxminddb.FromBytes(buffer)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	var record record

	err = db.Lookup(ip, &record)

	if err != nil {
		log.Fatal(err)
	}

	return record
}

func parseIPFromBilibiliAPI(ip net.IP) APIResponse {
	params := url.Values{}

	URL, err := url.Parse("https://api.live.bilibili.com/client/v1/Ip/getInfoNew")

	if err != nil {
		log.Fatal(err)
	}

	params.Set("ip", ip.String())

	URL.RawQuery = params.Encode()

	resp, err := http.Get(URL.String())

	if err != nil {
		log.Fatal(err)
	}

	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)
	var res APIResponse
	_ = json.Unmarshal(body, &res)
	return res
}

func printHop(hop traceroute.TracerouteHop) {
	addr := address(hop.Address)
	ip := net.ParseIP(addr)

	var info record
	var APIResponse APIResponse
	var err error

	if _, Prefix, _ := net.ParseCIDR("172.16.0.0/12"); Prefix.Contains(ip) {
		info, APIResponse, err = parseIPFromMap42(ip)
		if err != nil {
			info = parseIPFromMaxminddb(ip)
			APIResponse = parseIPFromBilibiliAPI(ip)
		}
	} else {
		info = parseIPFromMaxminddb(ip)
		APIResponse = parseIPFromBilibiliAPI(ip)
	}

	hostOrAddr := addr

	if hop.Host != "" {
		hostOrAddr = hop.Host
	}

	var ASN string

	if hop.Success {
		if info.ASN == 0 {
			ASN = "unknown"
		} else {
			ASN = fmt.Sprintf("AS%v", info.ASN)
		}
		if APIResponse.Code != 0 {
			fmt.Printf("%-3d %8v %15v %18v  %8v\n", hop.TTL, ASN, hostOrAddr, "("+addr+")", hop.ElapsedTime.Round(time.Microsecond))
		} else {
			fmt.Printf("%-3d %8v %15v %18v  %8v ", hop.TTL, ASN, hostOrAddr, "("+addr+")", hop.ElapsedTime.Round(time.Microsecond))

			var flag bool = false

			if APIResponse.Detail.Country != "" {
				flag = true
				fmt.Printf(APIResponse.Detail.Country)
			}
			if APIResponse.Detail.Province != "" {
				if flag {
					fmt.Printf(", ")
				}
				fmt.Print(APIResponse.Detail.Province)
				flag = true
			}
			if APIResponse.Detail.City != "" {
				if flag {
					fmt.Printf(", ")
				}
				fmt.Print(APIResponse.Detail.City)
				flag = true
			}

			if APIResponse.Detail.ISP != "" {
				if flag {
					fmt.Print(" ")
				}
				fmt.Print(APIResponse.Detail.ISP)
				flag = true
			}

			fmt.Printf("\n")
		}

	} else {
		fmt.Printf("%-3d *\n", hop.TTL)
	}

	if destIP.Equal(ip) {
		isFinish = true
	}

}

func address(address [4]byte) string {
	return fmt.Sprintf("%v.%v.%v.%v", address[0], address[1], address[2], address[3])
}

func main() {
	var m = flag.Int("m", traceroute.DEFAULT_MAX_HOPS, `Set the max time-to-live (max number of hops) used in outgoing probe packets (default is 64)`)
	var f = flag.Int("f", traceroute.DEFAULT_FIRST_HOP, `Set the first used time-to-live, e.g. the first hop (default is 1)`)
	var q = flag.Int("q", 1, `Set the number of probes per "ttl" to nqueries (default is one probe).`)

	flag.Parse()
	host := flag.Arg(0)
	options := traceroute.TracerouteOptions{}
	options.SetRetries(*q - 1)
	options.SetMaxHops(*m)
	options.SetFirstHop(*f)

	ipAddr, err := net.ResolveIPAddr("ip", host)
	if err != nil {
		return
	}

	fmt.Printf("traceroute to %v (%v), %v hops max, %v byte packets, using ICMP methods.\n", host, ipAddr, options.MaxHops(), options.PacketSize())

	c := make(chan traceroute.TracerouteHop, 0)
	d := make(chan bool)
	go func() {
		for {
			hop, ok := <-c
			if !ok {
				fmt.Println()
				return
			}
			printHop(hop)
			d <- true
		}
	}()

	destIP = ipAddr.IP

	_, err = traceroute.Traceroute(ipAddr.String(), &options, d, c)
	if err != nil {
		fmt.Printf("Error: %v", err)
	}
}