package main

import (
	"context"
	"errors"
	"fmt"
	"math"
	"net"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// Bounds on what a caller may ask: a check is a probe, never a wait anyone can make long.
const (
	DefaultTimeout = 3 * time.Second
	MostTimeout    = 30 * time.Second
)

// TCPResult is what netcheck_tcp answers.
type TCPResult struct {
	Host      string `json:"host"`
	Port      int    `json:"port"`
	Address   string `json:"address,omitempty"`
	Reachable bool   `json:"reachable"`
	ElapsedMS int64  `json:"elapsed_ms"`
	Error     string `json:"error,omitempty"`
}

// CheckTCP opens one TCP connection and closes it, sending nothing. A port that refuses or a host
// that does not answer is a result, not a failure of the tool; only a malformed question is.
func CheckTCP(host string, port int, timeout time.Duration) (TCPResult, error) {
	if port < 1 || port > 65535 {
		return TCPResult{}, fmt.Errorf("port %d is not a TCP port (1-65535)", port)
	}
	out := TCPResult{Host: host, Port: port}
	start := time.Now()
	conn, err := net.DialTimeout("tcp", net.JoinHostPort(host, strconv.Itoa(port)), timeout)
	out.ElapsedMS = time.Since(start).Milliseconds()
	if err != nil {
		out.Error = err.Error()
		return out, nil
	}
	out.Address = conn.RemoteAddr().String()
	out.Reachable = true
	_ = conn.Close()
	return out, nil
}

// DNSResult is what netcheck_dns answers.
type DNSResult struct {
	Name string `json:"name"`
	Type string `json:"type"`
	// Server and Protocol say which resolver was asked when one was named; empty, the machine's.
	Server   string `json:"server,omitempty"`
	Protocol string `json:"protocol,omitempty"`
	// Rcode is the named resolver's response code: NOERROR, NXDOMAIN, REFUSED, SERVFAIL…
	Rcode     string   `json:"rcode,omitempty"`
	Answers   []string `json:"answers"`
	ElapsedMS int64    `json:"elapsed_ms"`
	Error     string   `json:"error,omitempty"`
}

// DNSTypes are the record types netcheck_dns looks up.
var DNSTypes = []string{"A", "AAAA", "CNAME", "TXT", "MX"}

// CheckDNS looks a name up with the machine's resolver. Built without cgo, Go's own resolver reads
// the machine's /etc/resolv.conf and /etc/hosts, which is the resolver this machine's programs use.
// A name that does not resolve is a result with its error; an unknown type is refused.
func CheckDNS(name, kind string, timeout time.Duration) (DNSResult, error) {
	return AskDNS(name, kind, "", false, timeout)
}

// AskDNS is CheckDNS, or — with a server — the same question put to that resolver alone, over UDP or
// TCP, its hosts file and search domains left out (hq issue 348). No answer, a refusal and a response
// code other than NOERROR are results with their error; a malformed question is refused.
func AskDNS(name, kind, server string, overTCP bool, timeout time.Duration) (DNSResult, error) {
	kind = strings.ToUpper(strings.TrimSpace(kind))
	if kind == "" {
		kind = "A"
	}
	known := false
	for _, t := range DNSTypes {
		known = known || t == kind
	}
	if !known {
		return DNSResult{}, fmt.Errorf("type %q is not one netcheck_dns looks up (%s)", kind, strings.Join(DNSTypes, ", "))
	}
	out := DNSResult{Name: name, Type: kind, Answers: []string{}}
	if server != "" {
		if _, err := serverAddress(server); err != nil {
			return DNSResult{}, err
		}
		out.Server, out.Protocol = server, "udp"
		if overTCP {
			out.Protocol = "tcp"
		}
		start := time.Now()
		rcode, answers, err := askServer(server, name, kind, overTCP, timeout)
		out.ElapsedMS = time.Since(start).Milliseconds()
		if err == nil {
			out.Rcode = rcodeWords[rcode]
			if out.Rcode == "" {
				out.Rcode = "RCODE" + strconv.Itoa(rcode)
			}
			if rcode != 0 {
				out.Error = "the resolver answered " + out.Rcode
			}
		}
		out.Answers = append(out.Answers, answers...)
		if err != nil {
			out.Error = dnsWords(err, timeout)
		}
		if kind != "MX" {
			sort.Strings(out.Answers)
		}
		return out, nil
	} else if overTCP {
		return DNSResult{}, errors.New("tcp asks a named server; give server too")
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	r := net.DefaultResolver
	start := time.Now()
	var err error
	switch kind {
	case "A", "AAAA":
		network := "ip4"
		if kind == "AAAA" {
			network = "ip6"
		}
		var ips []net.IP
		if ips, err = r.LookupIP(ctx, network, name); err == nil {
			for _, ip := range ips {
				out.Answers = append(out.Answers, ip.String())
			}
		}
	case "CNAME":
		var cname string
		if cname, err = r.LookupCNAME(ctx, name); err == nil {
			out.Answers = append(out.Answers, cname)
		}
	case "TXT":
		var txts []string
		if txts, err = r.LookupTXT(ctx, name); err == nil {
			out.Answers = append(out.Answers, txts...)
		}
	case "MX":
		var mxs []*net.MX
		if mxs, err = r.LookupMX(ctx, name); err == nil {
			for _, mx := range mxs {
				out.Answers = append(out.Answers, fmt.Sprintf("%d %s", mx.Pref, mx.Host))
			}
		}
	}
	out.ElapsedMS = time.Since(start).Milliseconds()
	if err != nil {
		out.Error = err.Error()
	}
	if kind != "MX" {
		sort.Strings(out.Answers)
	}
	return out, nil
}

// text is a required string argument.
func text(args map[string]any, key string) (string, error) {
	s, _ := args[key].(string)
	s = strings.TrimSpace(s)
	if s == "" {
		return "", fmt.Errorf("%s is required", key)
	}
	return s, nil
}

// whole is an integer argument, given as a JSON number or a numeric string; fallback when absent.
func whole(args map[string]any, key string, fallback int) (int, error) {
	v, given := args[key]
	if !given || v == nil {
		if fallback == 0 {
			return 0, fmt.Errorf("%s is required", key)
		}
		return fallback, nil
	}
	switch n := v.(type) {
	case float64:
		if n != math.Trunc(n) {
			return 0, fmt.Errorf("%s must be a whole number, not %v", key, n)
		}
		return int(n), nil
	case string:
		i, err := strconv.Atoi(strings.TrimSpace(n))
		if err != nil {
			return 0, fmt.Errorf("%s must be a whole number, not %q", key, n)
		}
		return i, nil
	}
	return 0, errors.New(key + " must be a whole number")
}

// timeoutOf is timeout_ms, defaulted and bounded.
func timeoutOf(args map[string]any) (time.Duration, error) {
	ms, err := whole(args, "timeout_ms", int(DefaultTimeout/time.Millisecond))
	if err != nil {
		return 0, err
	}
	if ms < 1 {
		return 0, fmt.Errorf("timeout_ms must be at least 1, not %d", ms)
	}
	d := time.Duration(ms) * time.Millisecond
	if d > MostTimeout {
		d = MostTimeout
	}
	return d, nil
}

// dnsWords is a failed question in words: refused (nothing listens there), no answer in time, or what the
// system said.
func dnsWords(err error, timeout time.Duration) string {
	var ne net.Error
	switch {
	case errors.Is(err, syscall.ECONNREFUSED):
		return "connection refused: nothing listens on that address and port"
	case errors.As(err, &ne) && ne.Timeout():
		return fmt.Sprintf("no answer within %s", timeout)
	}
	return err.Error()
}