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

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"golang.org/x/text/encoding/charmap"
)

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

// maxAnswerSize caps the bytes read of one WHOIS answer: a real one is a
// few KiB, so one this large is from a server that does not stop sending.
const maxAnswerSize = 1 << 20

// Parsing constants.
const (
	lacnicDateFormatLen = 8 // Length of YYYYMMDD date format
)

// The classes of the two objects that describe an AS in an answer from
// RIPE, APNIC or AFRINIC, and the key of the aut-num object's line that
// gives the ID of the AS's organisation object.
const (
	autNumClass       = "aut-num"
	organisationClass = "organisation"
	orgKey            = "org"
)

// 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"

	// NIC.br registers Brazil's ASes under LACNIC.
	whoisServerNICbr = "whois.registro.br:43"
)

// nicbrHeader starts the copyright notice at the top of NIC.br's WHOIS
// answers. LACNIC's answer for an AS that NIC.br registers is NIC.br's, and
// carries the same notice:
//
//	% Copyright (c) Nic.br - Use of this data is governed by the Use and
//	% Privacy Policy at https://registro.br/upp . Distribution,
const nicbrHeader = "% Copyright (c) Nic.br"

// nicbrLimitReached is in the line that marks NIC.br's reduced answer to a
// client over its query limit. NIC.br publishes no limit. Its Use and
// Privacy Policy, https://registro.br/upp, says only that it restricts
// queries to stop the data being taken in bulk.
const nicbrLimitReached = "Query rate limit exceeded"

// NICbrSkipTime is how long the client stops asking NIC.br after NIC.br
// answers that the client is over its query limit. NIC.br publishes
// neither the limit nor how long it lasts. An hour outlasts a limit
// counted per minute or per hour. The ASes looked up meanwhile get
// NIC.br's fields when they are next refreshed.
const NICbrSkipTime = time.Hour

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

// noPersonalFlag keeps person and role objects out of an answer. The
// organisation object stays in it. RIPE counts the person and role objects
// it returns against a daily limit per IP address. Role objects for abuse
// contacts do not count. RIPE refuses an address for good that passes the
// limit too often:
// https://docs.db.ripe.net/Types-of-Queries/Access-Control-for-Queries/
//
// RIPE documents the flag here:
// https://docs.db.ripe.net/Types-of-Queries/Referenced-Objects-in-Query-Response/
// APNIC and AFRINIC list it in their answers to the query "help". AFRINIC
// sends readers to that answer for its full list of flags:
// https://afrinic.net/whois-getting-started
const noPersonalFlag = "--no-personal"

// The error lines with which RIPE's database software refuses a query:
// access denied to the client's IP address, and that address's limit
// reached. APNIC and AFRINIC run the same software. The lines are listed
// under "Access Errors" on this page, whose URL is split over two lines:
// https://docs.db.ripe.net/Appendices/
// Appendix-C--RIPE-Database-Query-Server-Response-Codes-and-Messages/
const (
	accessDeniedError = "%ERROR:201:"
	accessLimitError  = "%ERROR:202:"
)

// Errors LookupASN returns when a WHOIS answer has nothing to store.
var (
	// ErrNoReferral is returned when IANA refers the AS to no RIR, as it
	// does for a reserved, private-use or unallocated AS number.
	ErrNoReferral = errors.New("IANA refers the AS to no RIR")

	errNoASBlock = errors.New("IANA answered with neither a referral " +
		"nor an AS block")
	errNoFields       = errors.New("RIR answered with no WHOIS fields")
	errRefused        = errors.New("RIR refused the query")
	errUnknownServer  = errors.New("IANA refers the AS to a server no RIR runs")
	errAnswerTooLarge = errors.New("WHOIS answer reached the size cap")
)

// 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 {
	// DialContext opens the connection to a WHOIS server, given as
	// host:port. NewClient sets a net.Dialer's; tests replace it to reach
	// local servers instead.
	DialContext func(ctx context.Context, network, address string) (net.Conn, error)

	// Now returns the current time. NewClient sets time.Now; tests replace
	// it to move the clock.
	Now func() time.Time

	logger *slog.Logger

	// mu guards nicbrSkipUntil.
	mu sync.Mutex

	// nicbrSkipUntil is when the client next asks NIC.br, after NIC.br
	// answered that the client is over its query limit.
	nicbrSkipUntil time.Time
}

// NewClient creates a new WHOIS client that logs to logger.
func NewClient(logger *slog.Logger) *Client {
	dialer := &net.Dialer{Timeout: dialTimeout}

	return &Client{
		DialContext: dialer.DialContext,
		Now:         time.Now,
		logger:      logger,
	}
}

// LookupASN asks IANA which RIR holds the given ASN, asks that RIR, and
// returns what the RIR's answer says. When LACNIC's answer is NIC.br's, it
// returns what NIC.br's own server says instead, unless that query fails or
// is skipped, as askNICbr says.
//
// It returns ErrNoReferral when IANA refers the AS to no RIR, and another
// error when the query to IANA or to the RIR fails, when the RIR refuses
// it, when IANA's answer has neither a referral nor the AS's block, when
// IANA refers the AS to a server none of the five RIRs run, or when the
// answer carries none of ASNInfo's fields.
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, err := c.parseIANAReferral(ianaResp)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", query, err)
	}

	if whoisServer == "" {
		// IANA names the block of an AS it refers to no RIR. An answer
		// without that, such as an empty one, says nothing about the AS.
		if !strings.Contains(ianaResp, "as-block:") {
			return nil, fmt.Errorf("%s: %w", query, errNoASBlock)
		}

		return nil, fmt.Errorf("%s: %w", query, ErrNoReferral)
	}

	// Query the authoritative RIR
	rirResp, err := c.queryRIR(ctx, whoisServer, rirQuery(rir, query))
	if err != nil {
		return nil, fmt.Errorf("RIR query failed: %w", err)
	}

	// The fields never come from IANA's answer, which describes the block
	// the AS is in.
	info := c.parseResponse(asn, rir, rirResp)

	if rir == RIRLACNIC && strings.Contains(rirResp, nicbrHeader) {
		info = c.askNICbr(ctx, query, info)
	}

	if info.isEmpty() {
		return nil, fmt.Errorf("%s from %s: %w", query, whoisServer, errNoFields)
	}

	info.RawResponse = ianaResp + "\n" + info.RawResponse

	return info, nil
}

// askNICbr asks NIC.br's server about an AS whose answer from LACNIC is
// NIC.br's, and returns what NIC.br's answer says. LACNIC replaces each
// accented letter in such an answer with U+FFFD. NIC.br's own server keeps
// them, in Latin-1.
//
// askNICbr returns lacnicInfo, what LACNIC's answer says, and logs a
// warning when the query fails, when NIC.br sends the reduced answer for a
// client over its limit, or when NIC.br's answer carries none of ASNInfo's
// fields. After the reduced answer, it returns lacnicInfo without asking
// NIC.br until NICbrSkipTime has passed.
func (c *Client) askNICbr(
	ctx context.Context, asQuery string, lacnicInfo *ASNInfo,
) *ASNInfo {
	asn := lacnicInfo.ASN

	c.mu.Lock()
	skipping := c.Now().Before(c.nicbrSkipUntil)
	c.mu.Unlock()

	if skipping {
		return lacnicInfo
	}

	answer, err := c.query(ctx, whoisServerNICbr, asQuery)
	if err != nil {
		c.logger.Warn("NIC.br query failed, keeping LACNIC's fields",
			"asn", asn, "error", err)

		return lacnicInfo
	}

	if strings.Contains(answer, nicbrLimitReached) {
		c.mu.Lock()
		c.nicbrSkipUntil = c.Now().Add(NICbrSkipTime)
		c.mu.Unlock()

		c.logger.Warn("NIC.br answered that the client is over its query "+
			"limit, keeping LACNIC's fields and skipping NIC.br",
			"asn", asn, "skip_for", NICbrSkipTime)

		return lacnicInfo
	}

	info := c.parseResponse(asn, RIRLACNIC, answer)
	if info.isEmpty() {
		c.logger.Warn("NIC.br answered with no WHOIS fields, "+
			"keeping LACNIC's fields", "asn", asn)

		return lacnicInfo
	}

	return info
}

// isEmpty reports whether info holds none of the fields a WHOIS answer
// fills, as when the answer is only comment and error lines.
func (info *ASNInfo) isEmpty() bool {
	return info.ASName == "" && info.OrgName == "" && info.OrgID == "" &&
		info.Address == "" && info.CountryCode == "" &&
		info.AbuseEmail == "" && info.AbusePhone == "" &&
		info.TechEmail == "" && info.TechPhone == "" &&
		info.RegDate == nil && info.LastMod == nil
}

// rirQuery returns the query sent to rir's server for asQuery, an AS
// number such as "AS3333". RIPE, APNIC and AFRINIC are asked to keep person
// and role objects out of the answer.
func rirQuery(rir, asQuery string) string {
	switch rir {
	case RIRRIPE, RIRAPNIC, RIRAFRNIC:
		return noPersonalFlag + " " + asQuery
	default:
		return asQuery
	}
}

// queryRIR sends query to an RIR's server and returns the answer. An answer
// holding an error line with which the server refuses the query is an
// error, and the error holds that line.
func (c *Client) queryRIR(
	ctx context.Context, server, query string,
) (string, error) {
	answer, err := c.query(ctx, server, query)
	if err != nil {
		return "", err
	}

	for line := range strings.SplitSeq(answer, "\n") {
		if strings.HasPrefix(line, accessDeniedError) ||
			strings.HasPrefix(line, accessLimitError) {
			return "", fmt.Errorf("%w: %s", errRefused, line)
		}
	}

	return answer, nil
}

// query performs a raw WHOIS query to the specified server and returns the
// answer in UTF-8. An answer that reaches maxAnswerSize is an error.
func (c *Client) query(ctx context.Context, server, query string) (string, error) {
	conn, err := c.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

	// The reader ends at the cap; N, the bytes it has left, is 0 once the
	// answer has reached it.
	answer := &io.LimitedReader{R: conn, N: maxAnswerSize}

	scanner := bufio.NewScanner(answer)
	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)
	}

	if answer.N == 0 {
		return "", fmt.Errorf("read response: %w", errAnswerTooLarge)
	}

	return toUTF8(sb.String())
}

// toUTF8 returns a WHOIS answer in UTF-8. An answer that is valid UTF-8 is
// returned as it is. Any other answer is decoded as Latin-1, which RIPE,
// APNIC and NIC.br send. LACNIC, ARIN and AFRINIC send UTF-8. RIPE's
// database working group lists the five RIRs here:
// https://www.ripe.net/ripe/mail/archives/db-wg/2024-May/008075.html
func toUTF8(answer string) (string, error) {
	if utf8.ValidString(answer) {
		return answer, nil
	}

	decoded, err := charmap.ISO8859_1.NewDecoder().String(answer)
	if err != nil {
		return "", fmt.Errorf("decode Latin-1 answer: %w", err)
	}

	return decoded, nil
}

// parseIANAReferral extracts the RIR and WHOIS server from an IANA response.
// The server is "" when the response has no referral. A referral to a server
// that is none of the five RIRs' is an error, so a forged or broken answer
// cannot send the client to any other host or port.
func (c *Client) parseIANAReferral(response string) (string, string, error) {
	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, nil
			case strings.Contains(server, "ripe"):
				return RIRRIPE, whoisServerRIPE, nil
			case strings.Contains(server, "apnic"):
				return RIRAPNIC, whoisServerAPNIC, nil
			case strings.Contains(server, "lacnic"):
				return RIRLACNIC, whoisServerLACNIC, nil
			case strings.Contains(server, "afrinic"):
				return RIRAFRNIC, whoisServerAFRINIC, nil
			default:
				return "", "", fmt.Errorf("%w: %q", errUnknownServer, server)
			}
		}
	}

	return "", "", nil
}

// parseResponse extracts ASN information from a WHOIS response. Of an
// answer from RIPE, APNIC or AFRINIC it reads only the lines asObjectLines
// returns. Of an answer from ARIN or LACNIC it reads every line.
func (c *Client) parseResponse(asn int, rir, response string) *ASNInfo {
	info := &ASNInfo{
		ASN:         asn,
		RIR:         rir,
		RawResponse: response,
	}

	lines := strings.Split(response, "\n")
	if rir == RIRRIPE || rir == RIRAPNIC || rir == RIRAFRNIC {
		lines = asObjectLines(lines)
	}

	var addressLines []string

	for _, line := range lines {
		key, value, ok := splitField(line)
		if !ok || 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
}

// splitField splits a "key: value" line of a WHOIS answer into its key, in
// lower case, and its value. The third result is false for a line that
// holds no field: a blank line, a comment line or a line without a colon.
func splitField(line string) (string, string, bool) {
	line = strings.TrimSpace(line)
	if strings.HasPrefix(line, "%") || strings.HasPrefix(line, "#") {
		return "", "", false
	}

	key, value, ok := strings.Cut(line, ":")

	return strings.ToLower(strings.TrimSpace(key)), strings.TrimSpace(value), ok
}

// asObjectLines returns the lines of the aut-num object in an answer from
// RIPE, APNIC or AFRINIC, followed by the lines of the organisation object
// whose ID the aut-num object gives on its "org:" line. The answer's other
// objects are left out. RIPE and APNIC start their answer with the
// as-block object of the block the AS is in, whose dates are not the AS's.
//
// These answers are in RPSL. A blank line ends an object. An object's first
// line has the object's class as its key and the object's ID as its value,
// as in "organisation: ORG-APNI1-AP".
func asObjectLines(lines []string) []string {
	var autNum []string

	// The lines of each organisation object, by its ID.
	organisations := make(map[string][]string)

	// The class and ID of the object the line is in, "" between objects.
	var class, id string

	// The ID of the organisation object the aut-num object refers to.
	var orgID string

	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			class = ""

			continue
		}

		key, value, ok := splitField(line)
		if !ok {
			continue
		}

		if class == "" {
			class, id = key, value
		}

		switch class {
		case autNumClass:
			autNum = append(autNum, line)

			if key == orgKey {
				orgID = value
			}
		case organisationClass:
			organisations[id] = append(organisations[id], line)
		}
	}

	return append(autNum, organisations[orgID]...)
}

// 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 ""
}