package main

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

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

var (
	wlogger *log.Logger
)

func init() {
	wlogger = log.New(os.Stdout, "[ Whois ]", log.LstdFlags)
}

/*
func main() {
	pairs, err := findIPBlock("130.229.137.105", "192.88.99.1")
	if err != nil {
		panic(err)
	}
	log.Println(pairs)
}
*/

func findIPBlock(domains ...string) (pairs map[string]string, err error) {
	if VERBOSE {
		wlogger.Println("Querying for ip-blocks...")
	}
	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(map[string]string, len(lines)-2)
	for _, line := range lines[1 : len(lines)-1] {
		content := strings.Split(line, "|")
		ipaddr := strings.TrimSpace(content[1])
		ipb := strings.TrimSpace(content[2])
		pairs[ipaddr] = ipb
	}
	if VERBOSE {
		wlogger.Println("ip-blocks returned")
	}
	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
}