routewatch/internal/whois/whois.go

// Package whois provides WHOIS lookup functionality for ASN information.
package whois

import (
	"bufio"
	"context"
	"fmt"
	"net"
	"regexp"
	"strings"
	"time"
)

// Timeout constants for WHOIS queries.
const (
	dialTimeout  = 10 * time.Second
	readTimeout  = 30 * time.Second
	writeTimeout = 5 * time.Second
)

// Parsing constants.
const (
	keyValueParts       = 2 // Expected parts when splitting "key: value"
	lacnicDateFormatLen = 8 // Length of YYYYMMDD date format
)

// WHOIS server addresses.
const (
	whoisServerIANA    = "whois.iana.org:43"
	whoisServerARIN    = "whois.arin.net:43"
	whoisServerRIPE    = "whois.ripe.net:43"
	whoisServerAPNIC   = "whois.apnic.net:43"
	whoisServerLACNIC  = "whois.lacnic.net:43"
	whoisServerAFRINIC = "whois.afrinic.net:43"
)

// RIR identifiers.
const (
	RIRARIN   = "ARIN"
	RIRRIPE   = "RIPE"
	RIRAPNIC  = "APNIC"
	RIRLACNIC = "LACNIC"
	RIRAFRNIC = "AFRINIC"
)

// ASNInfo contains parsed WHOIS information for an ASN.
type ASNInfo struct {
	ASN         int
	ASName      string
	OrgName     string
	OrgID       string
	Address     string
	CountryCode string
	AbuseEmail  string
	AbusePhone  string
	TechEmail   string
	TechPhone   string
	RIR         string
	RegDate     *time.Time
	LastMod     *time.Time
	RawResponse string
}

// Client performs WHOIS lookups for ASNs.
type Client struct {
	// Dialer for creating connections (can be overridden for testing)
	dialer *net.Dialer
}

// NewClient creates a new WHOIS client.
func NewClient() *Client {
	return &Client{
		dialer: &net.Dialer{
			Timeout: dialTimeout,
		},
	}
}

// LookupASN queries WHOIS for the given ASN and returns parsed information.
func (c *Client) LookupASN(ctx context.Context, asn int) (*ASNInfo, error) {
	// Query IANA first to find the authoritative RIR
	query := fmt.Sprintf("AS%d", asn)

	ianaResp, err := c.query(ctx, whoisServerIANA, query)
	if err != nil {
		return nil, fmt.Errorf("IANA query failed: %w", err)
	}

	// Determine RIR from IANA response
	rir, whoisServer := c.parseIANAReferral(ianaResp)
	if whoisServer == "" {
		// No referral, try to parse what we have
		return c.parseResponse(asn, rir, ianaResp), nil
	}

	// Query the authoritative RIR
	rirResp, err := c.query(ctx, whoisServer, query)
	if err != nil {
		// Return partial data from IANA if RIR query fails
		info := c.parseResponse(asn, rir, ianaResp)
		info.RawResponse = fmt.Sprintf("%s\n--- RIR query failed: %v ---\n",
			ianaResp, err)

		return info, nil
	}

	// Combine responses and parse
	fullResponse := ianaResp + "\n" + rirResp
	info := c.parseResponse(asn, rir, fullResponse)
	info.RawResponse = fullResponse

	return info, nil
}

// query performs a raw WHOIS query to the specified server.
func (c *Client) query(ctx context.Context, server, query string) (string, error) {
	conn, err := c.dialer.DialContext(ctx, "tcp", server)
	if err != nil {
		return "", fmt.Errorf("dial %s: %w", server, err)
	}
	defer func() { _ = conn.Close() }()

	// Set deadlines
	err = conn.SetWriteDeadline(time.Now().Add(writeTimeout))
	if err != nil {
		return "", fmt.Errorf("set write deadline: %w", err)
	}

	// Send query
	_, err = fmt.Fprintf(conn, "%s\r\n", query)
	if err != nil {
		return "", fmt.Errorf("write query: %w", err)
	}

	// Read response
	err = conn.SetReadDeadline(time.Now().Add(readTimeout))
	if err != nil {
		return "", fmt.Errorf("set read deadline: %w", err)
	}

	var sb strings.Builder

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		sb.WriteString(scanner.Text())
		sb.WriteString("\n")
	}

	err = scanner.Err()
	if err != nil {
		return sb.String(), fmt.Errorf("read response: %w", err)
	}

	return sb.String(), nil
}

// parseIANAReferral extracts the RIR and WHOIS server from an IANA response.
func (c *Client) parseIANAReferral(response string) (string, string) {
	var rir string

	for line := range strings.SplitSeq(response, "\n") {
		line = strings.TrimSpace(line)

		// Look for "refer:" line
		if strings.HasPrefix(strings.ToLower(line), "refer:") {
			server := strings.TrimSpace(strings.TrimPrefix(line, "refer:"))
			server = strings.TrimSpace(strings.TrimPrefix(server, "Refer:"))

			switch {
			case strings.Contains(server, "arin"):
				return RIRARIN, whoisServerARIN
			case strings.Contains(server, "ripe"):
				return RIRRIPE, whoisServerRIPE
			case strings.Contains(server, "apnic"):
				return RIRAPNIC, whoisServerAPNIC
			case strings.Contains(server, "lacnic"):
				return RIRLACNIC, whoisServerLACNIC
			case strings.Contains(server, "afrinic"):
				return RIRAFRNIC, whoisServerAFRINIC
			default:
				// Unknown server, add port if missing
				if !strings.Contains(server, ":") {
					server += ":43"
				}

				return "", server
			}
		}

		// Also check organisation line for RIR hints
		if strings.HasPrefix(strings.ToLower(line), "organisation:") {
			orgRIR := rirInOrganisationLine(strings.ToLower(line))
			if orgRIR != "" {
				rir = orgRIR
			}
		}
	}

	return rir, ""
}

// rirInOrganisationLine returns the RIR a lowercased IANA "organisation:"
// line names, or "" when it names none.
func rirInOrganisationLine(org string) string {
	switch {
	case strings.Contains(org, "arin"):
		return RIRARIN
	case strings.Contains(org, "ripe"):
		return RIRRIPE
	case strings.Contains(org, "apnic"):
		return RIRAPNIC
	case strings.Contains(org, "lacnic"):
		return RIRLACNIC
	case strings.Contains(org, "afrinic"):
		return RIRAFRNIC
	default:
		return ""
	}
}

// parseResponse extracts ASN information from a WHOIS response.
func (c *Client) parseResponse(asn int, rir, response string) *ASNInfo {
	info := &ASNInfo{
		ASN:         asn,
		RIR:         rir,
		RawResponse: response,
	}

	lines := strings.Split(response, "\n")

	var addressLines []string

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "%") ||
			strings.HasPrefix(line, "#") {
			continue
		}

		// Split on first colon
		parts := strings.SplitN(line, ":", keyValueParts)
		if len(parts) != keyValueParts {
			continue
		}

		key := strings.TrimSpace(strings.ToLower(parts[0]))
		value := strings.TrimSpace(parts[1])

		if value == "" {
			continue
		}

		// Address (collect multiple lines)
		if key == "address" {
			addressLines = append(addressLines, value)

			continue
		}

		c.setResponseField(info, key, value)
	}

	// Combine address lines
	if len(addressLines) > 0 {
		info.Address = strings.Join(addressLines, "\n")
	}

	// Extract abuse email from comment lines (common in ARIN responses)
	if info.AbuseEmail == "" {
		info.AbuseEmail = c.extractAbuseEmail(response)
	}

	return info
}

// setResponseField stores the value of one "key: value" line of a WHOIS
// response in info, keeping the first value seen for each field.
func (c *Client) setResponseField(info *ASNInfo, key, value string) {
	switch key {
	// Country
	case "country":
		if info.CountryCode == "" && len(value) == 2 {
			info.CountryCode = strings.ToUpper(value)
		}

	// Registration dates
	case "regdate", "created":
		if info.RegDate == nil {
			info.RegDate = c.parseDate(value)
		}
	case "updated", "last-modified", "changed":
		if info.LastMod == nil {
			info.LastMod = c.parseDate(value)
		}

	default:
		field := textFieldForKey(info, key)
		if field != nil && *field == "" {
			*field = value
		}
	}
}

// textFieldForKey returns the text field of info that a WHOIS response key
// fills, or nil when the key fills none.
func textFieldForKey(info *ASNInfo, key string) *string {
	switch key {
	// AS Name (varies by RIR)
	case "asname", "as-name":
		return &info.ASName

	// Organization
	case "orgname", "org-name", "owner":
		return &info.OrgName
	case "orgid", "org-id", "org":
		return &info.OrgID

	// Abuse contact
	case "orgabuseemail", "abuse-mailbox":
		return &info.AbuseEmail
	case "orgabusephone":
		return &info.AbusePhone

	// Tech contact
	case "orgtechemail":
		return &info.TechEmail
	case "orgtechphone":
		return &info.TechPhone

	default:
		return nil
	}
}

// parseDate attempts to parse various date formats used in WHOIS responses.
func (c *Client) parseDate(value string) *time.Time {
	// Common formats
	formats := []string{
		"2006-01-02",
		"2006-01-02T15:04:05Z",
		"2006-01-02T15:04:05-07:00",
		"20060102",
		"02-Jan-2006",
	}

	// Clean up value
	value = strings.TrimSpace(value)
	// Handle "YYYYMMDD" format from LACNIC
	if len(value) == lacnicDateFormatLen {
		t, err := time.Parse("20060102", value)
		if err == nil {
			return &t
		}
	}

	for _, format := range formats {
		t, err := time.Parse(format, value)
		if err == nil {
			return &t
		}
	}

	return nil
}

// extractAbuseEmail extracts abuse email from response using regex.
func (c *Client) extractAbuseEmail(response string) string {
	// Look for "Abuse contact for 'AS...' is 'email@domain'"
	re := regexp.MustCompile(`[Aa]buse contact.*?is\s+['"]?([^\s'"]+@[^\s'"]+)['"]?`)
	if matches := re.FindStringSubmatch(response); len(matches) > 1 {
		return matches[1]
	}

	return ""
}