whoissdk/whois.go

package whois

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/proxy"
)

const (
	defWhoisServer = "whois.iana.org"
	defWhoisPort   = "43"
	defTimeout     = 30 * time.Second
	asnPrefix      = "AS"
)

const (
	StatusAddPeriod                = "AddPeriod"
	StatusAutoRenewPeriod          = "AutoRenewPeriod"
	StatusInActive                 = "Inactive"
	StatusOk                       = "Ok"
	StatusPendingDelete            = "PendingDelete"
	StatusPendingCreate            = "PendingCreate"
	StatusPendingRenew             = "PendingRenew"
	StatusPendingTransfer          = "PendingTransfer"
	StatusPendingRestore           = "PendingRestore"
	StatusPendingUpdate            = "PendingUpdate"
	StatusRedemptionPeriod         = "RedemptionPeriod"
	StatusRenewPeriod              = "RenewPeriod"
	StatusServerDeleteProhibited   = "ServerDeleteProhibited"
	StatusServerHold               = "ServerHold"
	StatusServerRenewProhibited    = "ServerRenewProhibited"
	StatusServerTransferProhibited = "ServerTransferProhibited"
	StatusServerUpdateProhibited   = "ServerUpdateProhibited"
	StatusTransferPeriod           = "TransferPeriod"
	StatusClientDeleteProhibited   = "ClientDeleteProhibited"
	StatusClientHold               = "ClientHold"
	StatusClientRenewProhibited    = "ClientRenewProhibited"
	StatusClientTransferProhibited = "ClientTransferProhibited"
	StatusClientUpdateProhibited   = "ClientUpdateProhibited"
)

/*
var statusMap = map[string]string{
	"addPeriod":                StatusAddPeriod,
	"autoRenewPeriod":          StatusAutoRenewPeriod,
	"inactive":                 StatusInActive,
	"ok":                       StatusOk,
	"pendingDelete":            StatusPendingDelete,
	"pendingCreate":            StatusPendingCreate,
	"pendingRenew":             StatusPendingRenew,
	"pendingTransfer":          StatusPendingTransfer,
	"pendingRestore":           StatusPendingRestore,
	"pendingUpdate":            StatusPendingUpdate,
	"redemptionPeriod":         StatusRedemptionPeriod,
	"renewPeriod":              StatusRenewPeriod,
	"serverDeleteProhibited":   StatusServerDeleteProhibited,
	"serverHold":               StatusServerHold,
	"serverRenewProhibited":    StatusServerRenewProhibited,
	"serverTransferProhibited": StatusServerTransferProhibited,
	"serverUpdateProhibited":   StatusServerUpdateProhibited,
	"transferPeriod":           StatusTransferPeriod,
	"clientDeleteProhibited":   StatusClientDeleteProhibited,
	"clientHold":               StatusClientHold,
	"clientRenewProhibited":    StatusClientRenewProhibited,
	"clientTransferProhibited": StatusClientTransferProhibited,
	"clientUpdateProhibited":   StatusClientUpdateProhibited,
}
*/

// DefaultClient is default whois client
var DefaultClient = NewClient()

var defaultWhoisMap = map[string]string{}

// Client is whois client
type Client struct {
	extCache        map[string]string
	mu              sync.Mutex
	dialer          proxy.Dialer
	timeout         time.Duration
	elapsed         time.Duration
	disableReferral bool
}
type PersonalInfo struct {
	FirstName string
	LastName  string
	Name      string
	Org       string
	Fax       string
	FaxExt    string
	Addr      string
	City      string
	State     string
	Country   string
	Zip       string
	Phone     string
	PhoneExt  string
	Email     string
}

type Result struct {
	exists          bool
	domain          string
	domainID        string
	rawData         string
	registar        string
	hasRegisterDate bool
	registerDate    time.Time
	hasUpdateDate   bool
	updateDate      time.Time
	hasExpireDate   bool
	expireDate      time.Time
	statusRaw       []string
	nsServers       []string
	nsIps           []string
	dnssec          string
	whoisSer        string
	ianaID          string
	registerInfo    PersonalInfo
	adminInfo       PersonalInfo
	techInfo        PersonalInfo
}

func (r Result) NsIps() []string {
	return r.nsIps
}

func (r Result) HasRegisterDate() bool {
	return r.hasRegisterDate
}

func (r Result) HasUpdateDate() bool {
	return r.hasUpdateDate
}

func (r Result) HasExpireDate() bool {
	return r.hasExpireDate
}

func (r Result) Exists() bool {
	return r.exists
}

func (r Result) Domain() string {
	return r.domain
}

func (r Result) DomainID() string {
	return r.domainID
}

func (r Result) RawData() string {
	return r.rawData
}

func (r Result) Registar() string {
	return r.registar
}

func (r Result) RegisterDate() time.Time {
	return r.registerDate
}

func (r Result) UpdateDate() time.Time {
	return r.updateDate
}

func (r Result) ExpireDate() time.Time {
	return r.expireDate
}

func (r Result) Status() []string {
	return r.statusRaw
}

func (r Result) NsServers() []string {
	return r.nsServers
}

func (r Result) Dnssec() string {
	return r.dnssec
}

func (r Result) WhoisSer() string {
	return r.whoisSer
}

func (r Result) IanaID() string {
	return r.ianaID
}

func (r Result) RegisterInfo() PersonalInfo {
	return r.registerInfo
}

func (r Result) AdminInfo() PersonalInfo {
	return r.adminInfo
}

func (r Result) TechInfo() PersonalInfo {
	return r.techInfo
}

// NewClient returns new whois client
func NewClient() *Client {
	return &Client{
		dialer: &net.Dialer{
			Timeout: defTimeout,
		},
		extCache: make(map[string]string),
		timeout:  defTimeout,
	}
}

func (c *Client) Whois(domain string, servers ...string) (Result, error) {
	if len(servers) == 0 {
		if v, ok := defaultWhoisMap[getExtension(domain)]; ok {
			servers = append(servers, v)
		}
	}
	data, err := c.whois(domain, servers...)
	if err != nil {
		return Result{}, err
	}
	return parse(domain, data)
}

func (c *Client) whois(domain string, servers ...string) (result string, err error) {
	start := time.Now()
	defer func() {
		result = strings.TrimSpace(result)
		if result != "" {
			result = fmt.Sprintf("%s\n\n%% Query time: %d msec\n%% WHEN: %s\n",
				result, time.Since(start).Milliseconds(), start.Format("Mon Jan 02 15:04:05 MST 2006"),
			)
		}
	}()

	domain = strings.Trim(strings.TrimSpace(domain), ".")
	if domain == "" {
		return "", errors.New("whois: domain is empty")
	}

	isASN := IsASN(domain)
	if isASN {
		if !strings.HasPrefix(strings.ToUpper(domain), asnPrefix) {
			domain = asnPrefix + domain
		}
	}

	if !strings.Contains(domain, ".") && !strings.Contains(domain, ":") && !isASN {
		return c.rawQuery(domain, defWhoisServer, defWhoisPort)
	}

	var server, port string
	if len(servers) > 0 && servers[0] != "" {
		server = strings.ToLower(servers[0])
		port = defWhoisPort
	} else {
		ext := getExtension(domain)
		if _, ok := c.extCache[ext]; !ok {
			result, err := c.rawQuery(ext, defWhoisServer, defWhoisPort)
			if err != nil {
				return "", fmt.Errorf("whois: query for whois server failed: %w", err)
			}
			server, port = getServer(result)
			if server == "" {
				return "", fmt.Errorf("%w: %s", errors.New("whois server not found"), domain)
			}
			c.mu.Lock()
			c.extCache[ext] = fmt.Sprintf("%s:%s", server, port)
			c.mu.Unlock()
		} else {
			v := strings.Split(c.extCache[ext], ":")
			server, port = v[0], v[1]
		}
	}

	result, err = c.rawQuery(domain, server, port)
	if err != nil {
		return
	}

	if c.disableReferral {
		return
	}

	refServer, refPort := getServer(result)
	if refServer == "" || refServer == server {
		return
	}

	data, err := c.rawQuery(domain, refServer, refPort)
	if err == nil {
		result += data
	}

	return
}

// SetDialer set query net dialer
func (c *Client) SetDialer(dialer proxy.Dialer) *Client {
	c.dialer = dialer
	return c
}

// SetTimeout set query timeout
func (c *Client) SetTimeout(timeout time.Duration) *Client {
	c.timeout = timeout
	return c
}

// rawQuery do raw query to the server
func (c *Client) rawQuery(domain, server, port string) (string, error) {
	c.elapsed = 0
	start := time.Now()

	if server == "whois.arin.net" {
		if IsASN(domain) {
			domain = "a + " + domain
		} else {
			domain = "n + " + domain
		}
	}

	conn, err := c.dialer.Dial("tcp", net.JoinHostPort(server, port))
	if err != nil {
		return "", fmt.Errorf("whois: connect to whois server failed: %w", err)
	}

	defer conn.Close()
	c.elapsed = time.Since(start)

	_ = conn.SetWriteDeadline(time.Now().Add(c.timeout - c.elapsed))
	_, err = conn.Write([]byte(domain + "\r\n"))
	if err != nil {
		return "", fmt.Errorf("whois: send to whois server failed: %w", err)
	}

	c.elapsed = time.Since(start)

	_ = conn.SetReadDeadline(time.Now().Add(c.timeout - c.elapsed))
	buffer, err := io.ReadAll(conn)
	if err != nil {
		return "", fmt.Errorf("whois: read from whois server failed: %w", err)
	}

	c.elapsed = time.Since(start)

	return string(buffer), nil
}

// getExtension returns extension of domain
func getExtension(domain string) string {
	ext := domain

	if net.ParseIP(domain) == nil {
		domains := strings.Split(domain, ".")
		ext = domains[len(domains)-1]
	}

	if strings.Contains(ext, "/") {
		ext = strings.Split(ext, "/")[0]
	}

	return ext
}

// getServer returns server from whois data
func getServer(data string) (string, string) {
	tokens := []string{
		"Registrar WHOIS Server: ",
		"whois: ",
		"ReferralServer: ",
		"refer: ",
	}
	for _, token := range tokens {
		start := strings.Index(data, token)
		if start != -1 {
			start += len(token)
			end := strings.Index(data[start:], "\n")
			server := strings.TrimSpace(data[start : start+end])
			server = strings.TrimPrefix(server, "http:")
			server = strings.TrimPrefix(server, "https:")
			server = strings.TrimPrefix(server, "whois:")
			server = strings.TrimPrefix(server, "rwhois:")
			server = strings.Trim(server, "/")
			port := defWhoisPort
			if strings.Contains(server, ":") {
				v := strings.Split(server, ":")
				server, port = v[0], v[1]
			}
			return server, port
		}
	}
	return "", ""
}

// IsASN returns if s is ASN
func IsASN(s string) bool {
	s = strings.ToUpper(s)

	s = strings.TrimPrefix(s, asnPrefix)
	_, err := strconv.Atoi(s)

	return err == nil
}