package main

import (
	"io/ioutil"
	"net"
	"strings"
	"time"
)

const (
	WHOIS_SERVER = /*"whois-servers.net"*/ "whois.cymru.com" /*"asn.shadowserver.org"*/
	WHOIS_PORT   = "43"
	TIMEOUT      = 30
	RETRIES      = 1
)

type asnipPair struct {
	asNum   string
	ipAdr   string
	ipBlock string
}

/*
func main() {
	pairs, err := findASAndIPBlock("109.105.104.100", "123.123.123.123")
	if err != nil {
		panic(err)
	}
	fmt.Println(pairs)
	pairs, err = findASAndIPBlock("123.123.123.123")
	if err != nil {
		panic(err)
	}
	fmt.Println(pairs)
}
*/

func findASAndIPBlock(domains ...string) (pairs []asnipPair, err error) {
	if len(domains) == 0 {
		return
	}

	var res string
	if len(domains) == 1 {
		res, err = whois(domains[0])
		if err != nil {
			return
		}
	} else {
		res, err = whoisBulk(domains)
		if err != nil {
			return
		}
	}

	lines := strings.Split(res, "\n")
	pairs = make([]asnipPair, 0, len(lines)-2)
	for ix, line := range lines[1 : len(lines)-1] {
		content := strings.Split(line, "|")
		as := strings.TrimSpace(content[0])
		ipb := strings.TrimSpace(content[2])
		pairs = append(pairs, asnipPair{
			asNum:   as,
			ipBlock: ipb,
			ipAdr:   domains[ix],
		})
	}
	return
}

func whois(domain string) (answer string, err error) {
	var conn net.Conn
	for i := 0; i < RETRIES; i++ {
		conn, err = net.DialTimeout("tcp", net.JoinHostPort("whois.cymru.com", WHOIS_PORT), time.Second*TIMEOUT)
		if err != nil {
			return
		}
	}
	defer conn.Close()

	//conn.Write([]byte(" -p " + domain + "\r\n"))
	conn.Write([]byte(" -p -o -u " + domain + "\r\n"))

	buff, err := ioutil.ReadAll(conn)
	if err != nil {
		return
	}

	answer = string(buff[:])
	return

}

func whoisBulk(domains []string) (answer string, err error) {
	var conn net.Conn
	for i := 0; i < RETRIES; i++ {
		conn, err = net.DialTimeout("tcp", net.JoinHostPort("whois.cymru.com", WHOIS_PORT), time.Second*TIMEOUT)
		if err != nil {
			return
		}
	}
	defer conn.Close()

	//conn.Write([]byte(" -p " + domain + "\r\n"))
	conn.Write([]byte(" -p -o -u begin\r\n"))
	for _, d := range domains {
		conn.Write([]byte(d + "\r\n"))
	}
	conn.Write([]byte("end\r\n"))

	buff, err := ioutil.ReadAll(conn)
	if err != nil {
		return
	}

	answer = string(buff[:])
	return
}