// Copyright Martin Dosch.
// Use of this source code is governed by the BSD-2-clause
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/tls"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"log"
	"math/big"
	"net"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"    // BSD-2-Clause
	"github.com/pborman/getopt/v2"    // BSD-3-Clause
	"golang.org/x/net/idna"           // BSD-3-Clause
	"salsa.debian.org/mdosch/xmppsrv" // BSD-2-Clause
)

const (
	version             = "0.6.2"
	nsBOSH              = "urn:xmpp:alt-connections:xbosh"
	nsC2SdTLS           = "urn:xmpp:alt-connections:tls"
	nsC2SQuic           = "urn:xmpp:alt-connections:quic"
	nsHTTPBind          = "http://jabber.org/protocol/httpbind"
	nsS2SdTLS           = "urn:xmpp:alt-connections:s2s-tls"
	nsS2SQuic           = "urn:xmpp:alt-connections:s2s-quic"
	nsS2SWebsocket      = "urn:xmpp:alt-connections:s2s-websocket"
	nsStartTLS          = "urn:ietf:params:xml:ns:xmpp-tls"
	nsWebsocket         = "urn:xmpp:alt-connections:websocket"
	nsXMPPFraming       = "urn:ietf:params:xml:ns:xmpp-framing"
	nsXMPPFramingServer = "urn:ietf:params:xml:ns:xmpp-framing-server"
)

// Created with https://github.com/miku/zek
type StreamError struct {
	XMLName xml.Name `xml:"stream"`
	Text    string   `xml:",chardata"`
	Version string   `xml:"version,attr"`
	Stream  string   `xml:"stream,attr"`
	Db      string   `xml:"db,attr"`
	Lang    string   `xml:"lang,attr"`
	From    string   `xml:"from,attr"`
	Xmlns   string   `xml:"xmlns,attr"`
	ID      string   `xml:"id,attr"`
	To      string   `xml:"to,attr"`
	Error   struct {
		Chardata string   `xml:",chardata"`
		Any      xml.Name `xml:",any"`
		Text     struct {
			Text  string `xml:",chardata"`
			Xmlns string `xml:"xmlns,attr"`
		} `xml:"text"`
	} `xml:"error"`
}

var (
	// statusOK will print [OK] in green color.
	statusOK = "Test: [\033[32mOK\033[00m]"
	// statusNOK will print [Not OK] in red color.
	statusNOK   = "Test: [\033[31mNot OK\033[00m]"
	colorGreen  = "\033[32m"
	colorRed    = "\033[31m"
	colorYellow = "\033[33m"
	colorReset  = "\033[00m"
)

func main() {
	// Define command line flags.
	flagHelp := getopt.BoolLong("help", 0, "Show help.")
	flagClient := getopt.BoolLong("client", 'c', "Show client SRV records.")
	flagServer := getopt.BoolLong("server", 's', "Show server SRV records.")
	flagVerbose := getopt.Bool('v', "Resolve IPs.")
	flagV4 := getopt.Bool('4', "Resolve IPv4.")
	flagV6 := getopt.Bool('6', "Resolve IPv6.")
	flagTest := getopt.Bool('t', "Test connection and certificates.")
	flagNoColor := getopt.BoolLong("no-color", 0, "Don't colorize output.")
	flagTLSVersion := getopt.IntLong("tls-version", 0, 12,
		"Minimal TLS version. 10 (TSLv1.0), 11 (TLSv1.1), 12 (TLSv1.2) or 13 (TLSv1.3).")
	flagVersion := getopt.BoolLong("version", 0, "Show version information.")
	flagFallback := getopt.BoolLong("fallback", 'f', "Additionally check fallback (Standard ports on A/AAAA records)"+
		" if no SRV records are provided.")
	flagForceFallback := getopt.BoolLong("force-fallback", 0, "Check fallback (Standard ports on A/AAAA records)"+
		" even if SRV records are provided.")
	flagTimeout := getopt.IntLong("timeout", 0, 60, "Connection timeout in seconds.")
	flagResolver := getopt.StringLong("resolver", 0, "", "Custom resolver e.g. \"1.1.1.1\" for common DNS"+
		" or \"5.1.66.255#dot.ffmuc.net\" for usage with \"--dot\".")
	flagDoT := getopt.BoolLong("dot", 0, "Use DNSoverTLS (DoT), see also \"--resolver\".")
	flagClientDTLS := getopt.ListLong("client-dtls", 0, "Add an arbitrary target for C2S with direct TLS in the form domain:port. Can be invoked several times.")
	flagClientSTLS := getopt.ListLong("client-stls", 0, "Add an arbitrary target for C2S with StartTLS in the form domain:port. Can be invoked several times.")
	flagServerDTLS := getopt.ListLong("server-dtls", 0, "Add an arbitrary target for S2S with direct TLS in the form domain:port. Can be invoked several times.")
	flagServerSTLS := getopt.ListLong("server-stls", 0, "Add an arbitrary target for S2S with StartTLS in the form domain:port. Can be invoked several times.")

	// Parse command line flags.
	getopt.Parse()

	// If requested, show help and quit.
	if *flagHelp {
		getopt.Usage()
		os.Exit(0)
	}

	// If requested, show version and quit.
	if *flagVersion {
		fmt.Println("Xmpp-dns", version)
		fmt.Println("Xmppsrv library version:", xmppsrv.Version)
		system := runtime.GOOS + "/" + runtime.GOARCH
		fmt.Println("System:", system, runtime.Version())
		fmt.Println("License: BSD-2-clause")
		os.Exit(0)
	}

	if *flagNoColor || runtime.GOOS == "windows" {
		statusOK = "Test: [OK]"
		statusNOK = "Test: [Not OK]"
		colorGreen = ""
		colorRed = ""
		colorYellow = ""
		colorReset = ""
	}

	// If connection test is required we'll also show IPs.
	if *flagTest && !*flagVerbose {
		*flagVerbose = true
	}

	// If either IPv4 or IPv6 is specified but the verbose flag
	// is not set, we'll just set it.
	if !*flagVerbose && (*flagV4 || *flagV6) {
		*flagVerbose = true
	}

	// If neither IPv4 nor Ipv6 is specified we'll use both.
	if !*flagV4 && !*flagV6 {
		*flagV4 = true
		*flagV6 = true
	}

	// If DoT is enabled a resolver must be set.
	if *flagDoT && *flagResolver == "" {
		log.Fatal("A resolver must be specified for DoT.")
	}

	// Read server from command line.
	server := getopt.Args()
	switch count := len(server); {
	case count == 0:
		log.Fatal("Please specify a server.")
	case count > 1:
		log.Fatal("Please specify only one server.")
	}
	// Convert server domain to ASCII.
	serverASCII, err := idna.ToASCII(server[0])
	if err != nil {
		log.Fatal(err)
	}

	// Configure DNS resolver
	c := xmppsrv.Config{
		Resolver: *flagResolver,
		DoT:      *flagDoT,
	}

	// Timeout
	timeout := time.Duration(*flagTimeout) * time.Second

	// Set TLS config
	var tlsConfig tls.Config
	tlsConfig.ServerName = serverASCII
	tlsConfig.InsecureSkipVerify = false
	switch *flagTLSVersion {
	case 10:
		tlsConfig.MinVersion = tls.VersionTLS10
	case 11:
		tlsConfig.MinVersion = tls.VersionTLS11
	case 12:
		tlsConfig.MinVersion = tls.VersionTLS12
	case 13:
		tlsConfig.MinVersion = tls.VersionTLS13
	default:
		fmt.Println("Unknown TLS version.")
		os.Exit(0)
	}

	// If neither client or server are chosen we default to showing both.
	if !*flagClient && !*flagServer {
		*flagClient = true
		*flagServer = true
	}

	if *flagClient {
		var fallback bool
		clientRecords, err := c.LookupClient(serverASCII)
		if err != nil && len(clientRecords) == 0 {
			fmt.Println(err)
			if *flagFallback && *flagTest {
				fallback = true
				fmt.Println("Trying fallback ports.")
				fmt.Println()
			}
		}
		if len(*flagClientDTLS) > 0 {
			for _, cdtls := range *flagClientDTLS {
				target := strings.SplitN(cdtls, ":", 2)
				if len(target) == 2 {
					port, err := strconv.ParseUint(target[1], 10, 16)
					if err != nil {
						log.Fatalf("client-dtls: failed to parse target (domain:port): %s: %v",
							cdtls, err)
					}
					cr := xmppsrv.SRV{
						Type:   "xmpps-client",
						Target: target[0],
						Port:   uint16(port),
					}
					clientRecords = append(clientRecords, cr)
				} else {
					log.Fatalf("client-dtls: failed to parse target (domain:port): %s\n", cdtls)
				}
			}
		}
		if len(*flagClientSTLS) > 0 {
			for _, cstls := range *flagClientSTLS {
				target := strings.SplitN(cstls, ":", 2)
				if len(target) == 2 {
					port, err := strconv.ParseUint(target[1], 10, 16)
					if err != nil {
						log.Fatalf("client-stls: failed to parse target (domain:port): %s: %v",
							cstls, err)
					}
					cr := xmppsrv.SRV{
						Type:   "xmpp-client",
						Target: target[0],
						Port:   uint16(port),
					}
					clientRecords = append(clientRecords, cr)
				} else {
					fmt.Printf("client-stls: failed to parse target (domain:port): %s\n", cstls)
				}
			}
		}
		if *flagForceFallback && !fallback && *flagTest {
			if len(clientRecords) > 0 {
				var standardPort bool
				for _, r := range clientRecords {
					if r.Type == "xmpp-client" &&
						(r.Target == server[0] ||
							r.Target == server[0]+".") &&
						r.Port == 5222 {
						standardPort = true
					}
				}
				if !standardPort {
					fallback = true
				}
			}
		}
		if fallback {
			cr := xmppsrv.SRV{
				Type:   "xmpp-client",
				Target: server[0],
				Port:   5222,
			}
			clientRecords = append(clientRecords, cr)
		}
		checkRecord(clientRecords, *flagVerbose, *flagV4, *flagV6, *flagTest,
			server[0], &tlsConfig, timeout)
	}

	if *flagServer {
		var fallback bool
		if *flagClient {
			fmt.Println()
		}
		serverRecords, err := c.LookupServer(serverASCII)
		if err != nil && len(serverRecords) == 0 {
			fmt.Println(err)
			if *flagFallback && *flagTest {
				fallback = true
				fmt.Println("Trying fallback ports.")
				fmt.Println()
			}
		}
		if len(*flagServerDTLS) > 0 {
			for _, sdtls := range *flagServerDTLS {
				target := strings.SplitN(sdtls, ":", 2)
				if len(target) == 2 {
					port, err := strconv.ParseUint(target[1], 10, 16)
					if err != nil {
						log.Fatalf("server-dtls: failed to parse target (domain:port): %s: %v",
							sdtls, err)
					}
					sr := xmppsrv.SRV{
						Type:   "xmpps-server",
						Target: target[0],
						Port:   uint16(port),
					}
					serverRecords = append(serverRecords, sr)
				} else {
					fmt.Printf("server-dtls: failed to parse target (domain:port): %s\n", sdtls)
				}
			}
		}
		if len(*flagServerSTLS) > 0 {
			for _, sstls := range *flagServerSTLS {
				target := strings.SplitN(sstls, ":", 2)
				if len(target) == 2 {
					port, err := strconv.ParseUint(target[1], 10, 16)
					if err != nil {
						log.Fatalf("server-stls: failed to parse target (domain:port): %s: %v",
							sstls, err)
					}
					sr := xmppsrv.SRV{
						Type:   "xmpp-server",
						Target: target[0],
						Port:   uint16(port),
					}
					serverRecords = append(serverRecords, sr)
				} else {
					fmt.Printf("server-stls: failed to parse target (domain:port): %s\n", sstls)
				}
			}
		}
		if *flagForceFallback && !fallback && *flagTest {
			if len(serverRecords) > 0 {
				var standardPort bool
				for _, r := range serverRecords {
					if r.Type == "xmpp-server" &&
						(r.Target == serverASCII ||
							r.Target == serverASCII+".") &&
						r.Port == 5269 {
						standardPort = true
					}
				}
				if !standardPort {
					fallback = true
				}
			}
		}
		if fallback {
			sr := xmppsrv.SRV{
				Type:   "xmpp-server",
				Target: serverASCII,
				Port:   5269,
			}
			serverRecords = append(serverRecords, sr)
		}
		checkRecord(serverRecords, *flagVerbose, *flagV4, *flagV6, *flagTest,
			server[0], &tlsConfig, timeout)
	}

	checkHostmeta(server[0], *flagVerbose, *flagV4, *flagV6,
		*flagClient, *flagServer, *flagTest, server[0], &tlsConfig, timeout, c)
	checkHostmeta2(server[0], *flagVerbose, *flagV4, *flagV6,
		*flagClient, *flagServer, *flagTest, server[0], &tlsConfig, timeout, c)
}

func checkRecord(records []xmppsrv.SRV, verbose bool, ipv4 bool, ipv6 bool, test bool,
	serverName string, tlsConfig *tls.Config, timeout time.Duration,
) {
	for count, record := range records {
		if count > 0 {
			fmt.Println()
		}
		printRecord(record)
		if verbose && record.Target != "." {
			printIP(record, ipv4, ipv6, test, serverName, tlsConfig, timeout)
		}
	}
}

func printRecord(record xmppsrv.SRV) {
	fmt.Println(record.Type, record.Target, record.Port)
	fmt.Print("Priority: ", record.Priority)
	fmt.Println(" Weight:", record.Weight)
}

func printIP(record xmppsrv.SRV, ipv4 bool, ipv6 bool, test bool,
	serverName string, tlsConfig *tls.Config, timeout time.Duration,
) {
	addresses, err := net.LookupIP(record.Target)
	switch {
	case err != nil:
		fmt.Println("No IP addresses found for", record.Target, err)
	case len(addresses) == 0:
		fmt.Println("No IP addresses found for", record.Target)
	default:
		var c net.Conn
		addresses = removeDuplicates(addresses)
		for _, address := range addresses {
			transport := "unset"
			if (address.To4() != nil) && ipv4 {
				transport = "tcp4"
				fmt.Println("IP:", address.To4())
			}
			if (address.To4() == nil) && ipv6 {
				transport = "tcp6"
				fmt.Println("IP:", address.To16())
			}
			if test && transport != "unset" {
				server := net.JoinHostPort(fmt.Sprint(address),
					fmt.Sprint(record.Port))
				c, err = connectionTest(server, transport, timeout)
				if err == nil {
					switch record.Type {
					case "xmpp-client":
						startTLS("client", serverName, c, tlsConfig, timeout)
						_ = c.Close()
					case "xmpps-client":
						tlsConfig.NextProtos = []string{"xmpp-client"}
						directTLS("client", serverName, c, tlsConfig, timeout)
						_ = c.Close()
					case "xmpp-server":
						startTLS("server", serverName, c, tlsConfig, timeout)
						_ = c.Close()
					case "xmpps-server":
						tlsConfig.NextProtos = []string{"xmpp-server"}
						directTLS("server", serverName, c, tlsConfig, timeout)
						_ = c.Close()
					default:
						_ = c.Close()
					}
				}
			}
		}
	}
}

func connectionTest(server string, transport string, timeout time.Duration) (net.Conn, error) {
	c, err := net.DialTimeout(transport, server, timeout)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return c, err
	}
	return c, err
}

func startTLS(recordType string, serverName string, c net.Conn, tlsConfig *tls.Config, timeout time.Duration) {
	// Created with https://github.com/miku/zek
	type Proceed struct {
		XMLName xml.Name `xml:"proceed"`
		Text    string   `xml:",chardata"`
		Xmlns   string   `xml:"xmlns,attr"`
	}
	// Created with https://github.com/miku/zek
	type Failure struct {
		XMLName xml.Name `xml:"failure"`
		Text    string   `xml:",chardata"`
		Xmlns   string   `xml:"xmlns,attr"`
	}
	var (
		serverProceed     Proceed
		serverFailure     Failure
		serverStreamError StreamError
	)
	var from string
	if recordType == "server" {
		from = " from='xmpp-dns.mdosch.de'"
	} else {
		from = ""
	}
	startStream := "<stream:stream xmlns:stream='http://etherx.jabber.org/streams' " +
		"xmlns='jabber:" + recordType + "'" + from + " to='" +
		serverName + "' version='1.0'>"
	_, err := c.Write([]byte(startStream))
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	var buf []byte
	bufChan := make(chan []byte)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		buf := make([]byte, 4096)
		_, err = c.Read(buf)
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				fmt.Println(statusNOK)
				if err.Error() == "EOF" {
					err = xml.Unmarshal(buf, &serverStreamError)
					if err == nil {
						fmt.Printf("stream error: %s: %s\n",
							serverStreamError.Error.Any.Local,
							serverStreamError.Error.Text.Text)
					} else {
						fmt.Println("Server closed connection.")
					}
				} else {
					fmt.Println(err)
				}
			}
		} else {
			bufChan <- buf
		}
	}()
	select {
	case buf = <-bufChan:
	case <-time.After(timeout):
	}
	if buf == nil {
		fmt.Println(statusNOK)
		fmt.Println("Timeout while waiting for server reply.")
		cancel()
		return
	}
	_ = xml.Unmarshal(buf, &serverFailure)
	if serverFailure.XMLName.Local == "failure" {
		fmt.Println(statusNOK)
		fmt.Println("Server sent failure.")
		return
	}
	_, err = fmt.Fprintf(c, "<starttls xmlns='%s'/>", nsStartTLS)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	for serverProceed.XMLName.Local != "proceed" &&
		serverProceed.Xmlns != nsStartTLS {
		_, err = c.Read(buf)
		if err != nil {
			fmt.Println(statusNOK)
			if err.Error() == "EOF" {
				err = xml.Unmarshal(buf, &serverStreamError)
				if err == nil {
					fmt.Printf("stream error: %s: %s\n",
						serverStreamError.Error.Any.Local,
						serverStreamError.Error.Text.Text)
				} else {
					fmt.Println("Server closed connection.")
				}
			} else {
				fmt.Println(err)
			}
			return
		}
		_ = xml.Unmarshal(buf, &serverProceed)
		_ = xml.Unmarshal(buf, &serverFailure)
		if serverFailure.XMLName.Local == "failure" {
			fmt.Println(statusNOK)
			fmt.Println("Server sent failure.")
			return
		}
	}
	d := tls.Client(c, tlsConfig)
	ctx2, cancel2 := context.WithTimeout(context.Background(), timeout)
	defer cancel2()
	err = d.HandshakeContext(ctx2)
	if err != nil {
		switch err.Error() {
		case "EOF":
			fmt.Println(statusNOK)
			fmt.Println("Server closed connection during handshake.")
		case "context deadline exceeded":
			fmt.Println(statusNOK)
			fmt.Println("Timeout during handshake.")
		default:
			fmt.Println(statusNOK)
			fmt.Println(err)
		}
	} else {
		checkCertExpiry(d)
		_ = d.Close()
	}
}

func directTLS(recordType string, serverName string, conn net.Conn, tlsConfig *tls.Config, timeout time.Duration) {
	var serverStreamError StreamError
	c := tls.Client(conn, tlsConfig)
	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	err := c.HandshakeContext(ctx)
	if err != nil {
		switch err.Error() {
		case "EOF":
			fmt.Println(statusNOK)
			fmt.Println("Server closed connection during handshake.")
		case "context deadline exceeded":
			fmt.Println(statusNOK)
			fmt.Println("Timeout during handshake.")
		default:
			fmt.Println(statusNOK)
			fmt.Println(err)
		}
		return
	} else {
		var from string
		if recordType == "server" {
			from = " from='xmpp-dns.mdosch.de'"
		} else {
			from = ""
		}
		startStream := "<stream:stream xmlns:stream='http://etherx.jabber.org/streams' " +
			"xmlns='jabber:" + recordType + "'" + from + " to='" +
			serverName + "' version='1.0'>"
		_, err := c.Write([]byte(startStream))
		if err != nil {
			fmt.Println(statusNOK)
			fmt.Println(err)
			return
		}
		buf := make([]byte, 4096)
		ctx3, cancel3 := context.WithCancel(context.Background())
		defer cancel3()
		go func() {
			for {
				_, err = c.Read(buf)
				if err != nil {
					fmt.Println(statusNOK)
					if err.Error() == "EOF" {
						err = xml.Unmarshal(buf, &serverStreamError)
						if err == nil {
							fmt.Println("Server sent policy-violation:",
								serverStreamError.Error.Text.Text)
						} else {
							fmt.Println("Server closed connection.")
						}
					} else {
						fmt.Println(err)
					}
					cancel3()
					return
				}
				if strings.Contains(strings.ToLower(string(buf)), "<stream:stream") &&
					(strings.Contains(strings.ToLower(string(buf)),
						"xmlns:stream='http://etherx.jabber.org/streams'") ||
						strings.Contains(strings.ToLower(string(buf)),
							`xmlns:stream="http://etherx.jabber.org/streams"`)) {
					checkCertExpiry(c)
					cancel3()
					break
				} else {
					continue
				}
			}
		}()
		select {
		case <-ctx3.Done():
		case <-time.After(timeout):
			fmt.Println(statusNOK)
			fmt.Println("Timeout during XMPP stream negotiation.")
		}
	}
}

func checkHostmeta(server string, verbose bool, ipv4 bool, ipv6 bool, c2s bool, s2s bool, test bool,
	serverName string, tlsConfig *tls.Config, timeout time.Duration, dnsConfig xmppsrv.Config,
) {
	type hostmetaXML struct {
		XMLName xml.Name `xml:"XRD"`
		Text    string   `xml:",chardata"`
		Xmlns   string   `xml:"xmlns,attr"`
		Link    []struct {
			Text string `xml:",chardata"`
			Rel  string `xml:"rel,attr"`
			Href string `xml:"href,attr"`
		} `xml:"Link"`
	}

	var hostmeta hostmetaXML

	httpClient, err := getHTTPClient(server, "443", ipv4, ipv6, timeout, tlsConfig, dnsConfig)
	if c2s && err == nil {
		resp, err := httpClient.Get("https://" + server + "/.well-known/host-meta")
		if err != nil {
			return
		}
		if resp.StatusCode == 404 {
			return
		}
		body, err := io.ReadAll(resp.Body)
		if err != nil {
			return
		}
		err = xml.Unmarshal(body, &hostmeta)
		if err != nil {
			return
		}
		var headlineShown bool
		for _, link := range hostmeta.Link {
			switch link.Rel {
			case nsWebsocket:
				if !headlineShown {
					fmt.Println("\nHost-meta:")
					headlineShown = true
				}
				fmt.Println("Websocket:", link.Href)
				if test {
					checkWebSocket(server, link.Href, tlsConfig, timeout, "c2s")
				}
			case nsBOSH:
				if !headlineShown {
					fmt.Println("\nHost-meta:")
					headlineShown = true
				}
				fmt.Println("BOSH:", link.Href)
				if test {
					checkBOSH(server, link.Href, ipv4, ipv6, tlsConfig,
						dnsConfig, timeout)
				}
			default:
				if isKnownLinkRel(link.Rel) {
					continue
				}
				if !headlineShown {
					fmt.Println("\nHost-meta:")
					headlineShown = true
				}
				fmt.Printf("Unknown connection method: %s\n", link.Rel)
			}
		}
		_ = resp.Body.Close()
	}
}

func checkHostmeta2(server string, verbose bool, ipv4 bool, ipv6 bool, c2s bool, s2s bool, test bool,
	serverName string, tlsConfig *tls.Config, timeout time.Duration, dnsConfig xmppsrv.Config,
) {
	type hostmetaJSON struct {
		XMPP struct {
			TTL                 int      `json:"ttl"`
			PublicKeyPinsSha256 []string `json:"public-key-pins-sha-256"`
		} `json:"xmpp"`
		Links []struct {
			Rel      string   `json:"rel"`
			Href     string   `json:"href,omitempty"`
			Ips      []string `json:"ips,omitempty"`
			Priority int      `json:"priority,omitempty"`
			Weight   int      `json:"weight,omitempty"`
			Sni      string   `json:"sni,omitempty"`
			Ech      string   `json:"ech,omitempty"`
			Port     int      `json:"port,omitempty"`
		} `json:"links"`
	}

	httpClient, err := getHTTPClient(server, "443", ipv4, ipv6, timeout, tlsConfig, dnsConfig)
	if err != nil {
		return
	}
	resp, err := httpClient.Get("https://" + server + "/.well-known/host-meta.json")
	if err != nil {
		return
	}
	if resp.StatusCode == 404 {
		return
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return
	}
	hostmeta2 := hostmetaJSON{}
	err = json.Unmarshal(body, &hostmeta2)
	if err != nil {
		return
	}
	if len(hostmeta2.Links) > 0 {
		var ipAddr string
		var headlineShown bool
		for _, link := range hostmeta2.Links {
			switch link.Rel {
			case nsWebsocket:
				if c2s {
					if !headlineShown {
						fmt.Println("\nHost-meta2:")
						headlineShown = true
					}
					fmt.Printf("C2S Websocket: %s\n", link.Href)
					if test {
						checkWebSocket(server, link.Href, tlsConfig, timeout, "c2s")
					}
				}
			case nsS2SWebsocket:
				if s2s {
					if !headlineShown {
						fmt.Println("\nHost-meta2:")
						headlineShown = true
					}
					fmt.Printf("S2S Websocket: %s\n", link.Href)
					if test {
						checkWebSocket(server, link.Href, tlsConfig, timeout, "s2s")
					}
				}
			case nsBOSH:
				if c2s {
					if !headlineShown {
						fmt.Println("\nHost-meta2:")
						headlineShown = true
					}
					fmt.Printf("BOSH: %s\n", link.Href)
					if test {
						checkBOSH(server, link.Href, ipv4, ipv6, tlsConfig,
							dnsConfig, timeout)
					}
				}
			case nsC2SdTLS:
				if c2s {
					if !headlineShown {
						fmt.Println("\nHost-meta2:")
						headlineShown = true
					}
					fmt.Printf("xmpps-client %s %d\n", link.Sni, link.Port)
					fmt.Printf("Priority: %d Weight %d\n", link.Priority, link.Weight)
					if verbose {
						for _, ip := range link.Ips {
							transport := "unset"
							address := net.ParseIP(ip)
							if (address.To4() != nil) && ipv4 {
								transport = "tcp4"
								ipAddr = address.To4().String()
							}
							if (address.To4() == nil) && ipv6 {
								transport = "tcp6"
								ipAddr = address.To16().String()
							}
							fmt.Println("IP:", ipAddr)
							if test && transport != "unset" {
								var c net.Conn
								server := net.JoinHostPort(fmt.Sprint(address),
									fmt.Sprint(link.Port))
								c, err = connectionTest(server, transport, timeout)
								if err != nil {
									fmt.Println(statusNOK)
									continue
								}
								tlsConfig.NextProtos = []string{"xmpp-client"}
								directTLS("client", serverName, c, tlsConfig, timeout)
								_ = c.Close()
							}
						}
					}
				}
			case nsS2SdTLS:
				if s2s {
					if !headlineShown {
						fmt.Println("\nHost-meta2:")
						headlineShown = true
					}
					fmt.Printf("xmpps-server %s %d\n", link.Sni, link.Port)
					fmt.Printf("Priority: %d Weight %d\n", link.Priority, link.Weight)
					if verbose {
						for _, ip := range link.Ips {
							transport := "unset"
							address := net.ParseIP(ip)
							if (address.To4() != nil) && ipv4 {
								transport = "tcp4"
								ipAddr = address.To4().String()
							}
							if (address.To4() == nil) && ipv6 {
								transport = "tcp6"
								ipAddr = address.To16().String()
							}
							fmt.Println("IP:", ipAddr)
							if test && transport != "unset" {
								var c net.Conn
								server := net.JoinHostPort(fmt.Sprint(address),
									fmt.Sprint(link.Port))
								c, err = connectionTest(server, transport, timeout)
								if err != nil {
									fmt.Println(statusNOK)
									continue
								}
								tlsConfig.NextProtos = []string{"xmpp-server"}
								directTLS("client", serverName, c, tlsConfig, timeout)
								_ = c.Close()
							}
						}
					}
				}
			case nsC2SQuic:
				if c2s {
					if !headlineShown {
						fmt.Println("\nHost-meta2:")
						headlineShown = true
					}
					fmt.Printf("C2S QUIC: %s %d\n", link.Sni, link.Port)
					if verbose {
						for _, ip := range link.Ips {
							address := net.ParseIP(ip)
							if (address.To4() != nil) && ipv4 {
								// connType = "tcp4"
								ipAddr = address.To4().String()
							}
							if (address.To4() == nil) && ipv6 {
								// connType = "tcp6"
								ipAddr = address.To16().String()
							}
							fmt.Println("IP:", ipAddr)
							if test {
								fmt.Println("No QUIC test yet")
							}
						}
					}
				}
			case nsS2SQuic:
				if s2s {
					if !headlineShown {
						fmt.Println("\nHost-meta2:")
						headlineShown = true
					}
					fmt.Printf("S2S QUIC: %s %d\n", link.Sni, link.Port)
					if verbose {
						for _, ip := range link.Ips {
							address := net.ParseIP(ip)
							if (address.To4() != nil) && ipv4 {
								// connType = "tcp4"
								ipAddr = address.To4().String()
							}
							if (address.To4() == nil) && ipv6 {
								// connType = "tcp6"
								ipAddr = address.To16().String()
							}
							fmt.Println("IP:", ipAddr)
							if test {
								fmt.Println("No QUIC test yet")
							}
						}
					}
				}
			default:
				if isKnownLinkRel(link.Rel) {
					continue
				}
				if !headlineShown {
					fmt.Println("\nHost-meta2:")
					headlineShown = true
				}
				fmt.Printf("Unknown connection method: %s\n", link.Rel)
			}
		}
	}
	_ = resp.Body.Close()
}

func checkBOSH(server string, target string, ipv4 bool, ipv6 bool, tlsConfig *tls.Config,
	srvConfig xmppsrv.Config, timeout time.Duration,
) {
	type boshFailXML struct {
		XMLName   xml.Name `xml:"body"`
		Text      string   `xml:",chardata"`
		Type      string   `xml:"type,attr"`
		Condition string   `xml:"condition,attr"`
		Xmlns     string   `xml:"xmlns,attr"`
		Stream    string   `xml:"stream,attr"`
	}
	type boshOpenXML struct {
		XMLName    xml.Name `xml:"body"`
		Text       string   `xml:",chardata"`
		Wait       string   `xml:"wait,attr"`
		Inactivity string   `xml:"inactivity,attr"`
		Polling    string   `xml:"polling,attr"`
		Requests   string   `xml:"requests,attr"`
		Hold       string   `xml:"hold,attr"`
		Ack        string   `xml:"ack,attr"`
		Accept     string   `xml:"accept,attr"`
		Maxpause   string   `xml:"maxpause,attr"`
		Sid        string   `xml:"sid,attr"`
		Charsets   string   `xml:"charsets,attr"`
		Ver        string   `xml:"ver,attr"`
		From       string   `xml:"from,attr"`
		Xmlns      string   `xml:"xmlns,attr"`
	}

	var boshFail boshFailXML
	var boshOpen boshOpenXML

	boshURI, err := url.ParseRequestURI(target)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	bodyXML := `<body content='text/xml; charset=utf-8' hold='1' rid='` + getID() +
		`' to='` + server + `' ver='1.6' wait='0' ack='1' xml:lang='en'` +
		` xmlns='` + nsHTTPBind + `'/>`
	body := []byte(bodyXML)
	req, err := http.NewRequest("POST", target, bytes.NewBuffer(body))
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	tlsConfig.NextProtos = []string{"http/1.1", "h2"}
	tlsConfig.ServerName = boshURI.Hostname()
	httpClient, err := getHTTPClient(boshURI.Hostname(), boshURI.Port(), ipv4, ipv6, timeout, tlsConfig, srvConfig)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	resp, err := httpClient.Do(req)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return
	}
	err = xml.Unmarshal(respBody, &boshFail)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	if boshFail.Type == "terminate" {
		fmt.Println(statusNOK)
		fmt.Printf("Server terminated: %s\n", boshFail.Condition)
		return
	}
	err = xml.Unmarshal(respBody, &boshOpen)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	if boshOpen.From != server {
		fmt.Println(statusNOK)
		fmt.Printf("Server identifies as %s instead of %s.\n",
			boshOpen.From, server)
		return
	}
	if resp.TLS != nil {
		fmt.Print(statusOK)
		printTLSVersion(resp.TLS.Version)
	} else {
		fmt.Printf("%s (%sno TLS%s)", statusNOK, colorRed, colorReset)
	}
	fmt.Println()
	_ = resp.Body.Close()
}

func checkWebSocket(server string, target string, tlsConfig *tls.Config, timeout time.Duration, targetType string) {
	type wsOpenXML struct {
		XMLName xml.Name `xml:"open"`
		Text    string   `xml:",chardata"`
		Xmlns   string   `xml:"xmlns,attr"`
		ID      string   `xml:"id,attr"`
		Version string   `xml:"version,attr"`
		Lang    string   `xml:"lang,attr"`
		From    string   `xml:"from,attr"`
	}
	var wsOpen wsOpenXML
	var framing, from string
	var dialer websocket.Dialer
	switch targetType {
	case "c2s":
		framing = nsXMPPFraming
	case "s2s":
		framing = nsXMPPFramingServer
		from = " from='xmpp-dns.mdosch.de'"
	}
	wsURI, err := url.ParseRequestURI(target)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	origURI, err := url.ParseRequestURI("https://" + server)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	headers := http.Header{}
	headers.Add("Origin", origURI.String())
	wsTLSConfig := &tls.Config{
		MinVersion:         tlsConfig.MinVersion,
		ServerName:         wsURI.Hostname(),
		InsecureSkipVerify: false,
	}
	// Websocket connections currently fail when using HTTP/2, but also none of the
	// alternative libraries support it, so currently it is disabled.
	// TODO: Switch library, once one supports it.
	// See:
	// https://github.com/coder/websocket/issues/4
	// https://github.com/gorilla/websocket/issues/417
	// wsTLSConfig.NextProtos = []string{"http/1.1", "h2"}
	wsTLSConfig.NextProtos = []string{"http/1.1"}
	dialer.TLSClientConfig = wsTLSConfig
	dialer.Subprotocols = append(dialer.Subprotocols, "xmpp")
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	wsConn, resp, err := dialer.DialContext(ctx, wsURI.String(), headers)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	err = wsConn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(
		`<open xmlns='%s' to='%s'%s version='1.0' />`, framing, server, from)))
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	_, buf, err := wsConn.ReadMessage()
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	err = xml.Unmarshal(buf, &wsOpen)
	if err != nil {
		fmt.Println(statusNOK)
		fmt.Println(err)
		return
	}
	if wsOpen.From != server {
		_ = wsConn.Close()
		fmt.Println(statusNOK)
		fmt.Printf("Server identifies as %s instead of %s.\n",
			wsOpen.From, server)
		return
	}
	if resp.TLS != nil {
		fmt.Print(statusOK)
		printTLSVersion(resp.TLS.Version)
		fmt.Println()
	} else {
		state := getTLS(wsConn)
		if state != nil {
			fmt.Print(statusOK)
			printTLSVersion(state.Version)
			fmt.Println()
		} else {
			fmt.Printf("%s (%sno TLS%s)", statusNOK, colorRed, colorReset)
			fmt.Println()
		}
	}
	err = wsConn.Close()
	if err != nil {
		fmt.Println(err)
	}
}

func getHTTPClient(server string, port string, ipv4 bool, ipv6 bool, timeout time.Duration,
	tlsConfig *tls.Config, dnsConfig xmppsrv.Config,
) (*http.Client, error) {
	var ipAddr, connType string
	addresses, err := dnsConfig.GetIP(server)
	switch {
	case err != nil:
		return &http.Client{}, fmt.Errorf("no IP addresses found for %s: %v", server, err)
	case len(addresses) == 0:
		return &http.Client{}, fmt.Errorf("no IP addresses found for %s", server)
	default:
		addresses = removeDuplicates(addresses)
		connType = "tcp"
		for _, address := range addresses {
			if (address.To4() != nil) && ipv4 {
				connType = "tcp4"
				ipAddr = address.To4().String()
			}
			if (address.To4() == nil) && ipv6 {
				connType = "tcp6"
				ipAddr = address.To16().String()
			}
		}
	}
	httpTLSConfig := &tls.Config{
		MinVersion:         tlsConfig.MinVersion,
		ServerName:         server,
		InsecureSkipVerify: false,
	}
	if port == "" {
		port = "443"
	}
	target := net.JoinHostPort(ipAddr, port)
	httpTransport := &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			dialer := net.Dialer{}
			ctx2, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()
			return dialer.DialContext(ctx2, connType, target)
		},
		IdleConnTimeout:     timeout,
		TLSClientConfig:     httpTLSConfig,
		TLSHandshakeTimeout: timeout,
	}
	httpClient := &http.Client{
		Transport:     httpTransport,
		CheckRedirect: nil,
	}
	return httpClient, nil
}

func checkCertExpiry(c *tls.Conn) {
	expiry := c.ConnectionState().PeerCertificates[0].NotAfter
	start := c.ConnectionState().PeerCertificates[0].NotBefore
	now := time.Now()
	switch {
	case time.Until(expiry).Hours() < 336:
		fmt.Printf("Test: [%sWARN%s]", colorYellow, colorReset)
		printTLSVersion(c.ConnectionState().Version)
		expiryDays := int(time.Until(expiry).Hours())/24 + 1
		dayString := fmt.Sprintf("%d days", expiryDays)
		if expiryDays == 1 {
			dayString = "one day"
		}
		fmt.Printf(" (Cert expires in less than %s at %s)\n",
			dayString, expiry.Truncate(time.Second).String())
	case now.Before(expiry) && now.After(start):
		fmt.Print(statusOK)
		printTLSVersion(c.ConnectionState().Version)
		fmt.Println()
	default:
		fmt.Print(statusNOK)
		printTLSVersion(c.ConnectionState().Version)
		fmt.Printf(" (valid from %s to %s)", start, expiry)
	}
}

func printTLSVersion(tv uint16) {
	switch tv {
	case tls.VersionTLS13:
		fmt.Printf(" (%s%s%s)", colorGreen, tls.VersionName(tv), colorReset)
	case tls.VersionTLS12:
		fmt.Printf(" (%s%s%s)", colorYellow, tls.VersionName(tv), colorReset)
	default:
		fmt.Printf(" (%s%s%s)", colorRed, tls.VersionName(tv), colorReset)
	}
}

func removeDuplicates(addresses []net.IP) []net.IP {
	var addressList []net.IP
	for _, ip := range addresses {
		if !containsIP(addressList, ip) {
			addressList = append(addressList, ip)
		}
	}
	return addressList
}

func containsIP(addressList []net.IP, ip net.IP) bool {
	for _, r := range addressList {
		if ip.Equal(r) {
			return true
		}
	}
	return false
}

func getID() string {
	idRunes := []rune("0123456789")
	max := big.NewInt(int64(len(idRunes)))
	id := make([]rune, 10)
	for i := range id {
		randInt, err := rand.Int(rand.Reader, max)
		if err != nil {
			log.Fatal(err)
		}
		id[i] = idRunes[randInt.Int64()]
	}
	return string(id)
}

func getTLS(c any) *tls.ConnectionState {
	for c != nil {
		if tc, ok := c.(*tls.Conn); ok {
			state := tc.ConnectionState()
			return &state
		}
		wc, ok := c.(interface{ NetConn() net.Conn })
		if !ok {
			break
		}
		c = wc.NetConn()
	}
	return nil
}
