// Copyright 2021 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 (
	"crypto/tls"
	"encoding/xml"
	"fmt"
	"log"
	"net"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/pborman/getopt/v2"    // BSD-3-Clause
	"salsa.debian.org/mdosch/xmppsrv" // BSD-2-Clause
)

const (
	version = "0.3.4"
)

var (
	// statusOK will print [OK] in green color.
	statusOK = "[\033[32mOK\033[00m]"
	// statusNOK will print [Not OK] in red color.
	statusNOK = "[\033[31mNot OK\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', "Check fallback (Standard ports on A/AAAA records)"+
		" if no 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\".")

	// 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("License: BSD-2-clause")
		os.Exit(0)
	}

	if *flagNoColor || runtime.GOOS == "windows" {
		statusOK = "[OK]"
		statusNOK = "[Not OK]"
	}

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

	// If verbose output (showing IPs) is requested but neither IPv4 nor Ipv6
	// is specified we'll show both.
	if *flagVerbose && !*flagV4 && !*flagV6 {
		*flagV4 = true
		*flagV6 = 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 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.")
	}

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

	// Lookup CNAME
	target, err := c.LookupCNAME(server[0])
	if err != nil {
		log.Fatal(err)
	}

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

	// Set TLS config
	var tlsConfig tls.Config
	tlsConfig.ServerName = server[0]
	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 {
		clientRecords, err := c.LookupClient(target)
		if err != nil && len(clientRecords) == 0 {
			fmt.Println(err)
			if *flagFallback && *flagTest {
				fmt.Println("Trying fallback ports.")
				fmt.Println()
				clientRecords = []xmppsrv.SRV{
					{
						Type:   "xmpp-client",
						Target: target,
						Port:   5222,
					},
					{
						Type:   "xmpps-client",
						Target: target,
						Port:   5223,
					},
					{
						Type:   "xmpps-client",
						Target: target,
						Port:   443,
					},
				}
			}
		}
		checkRecord(clientRecords, *flagVerbose, *flagV4, *flagV6, *flagTest,
			&tlsConfig, timeout)
	}

	if *flagServer {
		if *flagClient {
			fmt.Println()
		}
		serverRecords, err := c.LookupServer(target)
		if err != nil && len(serverRecords) == 0 {
			fmt.Println(err)
			if *flagFallback && *flagTest {
				fmt.Println("Trying fallback ports.")
				fmt.Println()
				serverRecords = []xmppsrv.SRV{
					{
						Type:   "xmpp-server",
						Target: target,
						Port:   5269,
					},
					{
						Type:   "xmpps-server",
						Target: target,
						Port:   5270,
					},
				}
			}
		}
		checkRecord(serverRecords, *flagVerbose, *flagV4, *flagV6, *flagTest,
			&tlsConfig, timeout)
	}
}

func checkRecord(records []xmppsrv.SRV, verbose bool, ipv4 bool, ipv6 bool, test bool,
	tlsConfig *tls.Config, timeout time.Duration) {
	for count, record := range records {
		if count > 0 {
			fmt.Println()
		}
		printRecord(record)
		if verbose {
			printIP(record, ipv4, ipv6, test, 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,
	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", c, tlsConfig)
						c.Close()
					case "xmpps-client":
						tlsConfig.NextProtos = []string{"xmpp-client"}
						directTLS("client", c, tlsConfig)
						c.Close()
					case "xmpp-server":
						startTLS("server", c, tlsConfig)
						c.Close()
					case "xmpps-server":
						tlsConfig.NextProtos = []string{"xmpp-server"}
						directTLS("server", c, tlsConfig)
						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("Connection:", statusNOK)
		fmt.Println(err)
		return c, err
	}
	fmt.Println("Connection:", statusOK)
	return c, err
}

func startTLS(recordType string, c net.Conn, tlsConfig *tls.Config) {
	// 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
	)
	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='" +
		tlsConfig.ServerName + "' version='1.0'>\n"
	_, err := c.Write([]byte(startStream))
	if err != nil {
		fmt.Println("StartTLS:", statusNOK)
		fmt.Println(err)
		return
	}
	buf := make([]byte, 4096)
	_, err = c.Read(buf)
	if err != nil {
		fmt.Println("StartTLS:", statusNOK)
		if err.Error() == "EOF" {
			fmt.Println("Server sent EOF.")
		} else {
			fmt.Println(err)
		}
		return
	}
	_ = xml.Unmarshal(buf, &serverFailure)
	if serverFailure.XMLName.Local == "failure" {
		fmt.Println("StartTLS:", statusNOK)
		fmt.Println("Server sent failure.")
		return
	}
	_, err = c.Write([]byte("<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>\n"))
	if err != nil {
		fmt.Println("StartTLS:", statusNOK)
		fmt.Println(err)
		return
	}
	for !(serverProceed.XMLName.Local == "proceed" &&
		serverProceed.Xmlns == "urn:ietf:params:xml:ns:xmpp-tls") {
		_, err = c.Read(buf)
		if err != nil {
			fmt.Println("StartTLS:", statusNOK)
			if err.Error() == "EOF" {
				fmt.Println("Server sent EOF.")
			} else {
				fmt.Println(err)
			}
			return
		}
		_ = xml.Unmarshal(buf, &serverProceed)
		_ = xml.Unmarshal(buf, &serverFailure)
		if serverFailure.XMLName.Local == "failure" {
			fmt.Println("StartTLS:", statusNOK)
			fmt.Println("Server sent failure.")
			return
		}
	}
	d := tls.Client(c, tlsConfig)
	err = d.Handshake()
	if err != nil {
		if err.Error() == "EOF" {
			fmt.Println("StartTLS:", statusNOK)
			fmt.Println("Received EOF during handshake.")
		} else {
			fmt.Println("StartTLS:", statusNOK)
			fmt.Println(err)
		}
	} else {
		fmt.Println("StartTLS:", statusOK)
		checkCertExpiry(d)
		d.Close()
	}
}

func directTLS(recordType string, conn net.Conn, tlsConfig *tls.Config) {
	c := tls.Client(conn, tlsConfig)
	err := c.Handshake()
	if err != nil {
		if err.Error() == "EOF" {
			fmt.Println("TLS:", statusNOK)
			fmt.Println("Received EOF during handshake.")
		} else {
			fmt.Println("TLS:", statusNOK)
			fmt.Println(err)
		}
	} 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='" +
			tlsConfig.ServerName + "' version='1.0'>\n"
		_, err := c.Write([]byte(startStream))
		if err != nil {
			fmt.Println("TLS:", statusNOK)
			fmt.Println(err)
			return
		}
		buf := make([]byte, 4096)
		_, err = c.Read(buf)
		if err != nil {
			fmt.Println("TLS:", statusNOK)
			if err.Error() == "EOF" {
				fmt.Println("Server sent EOF.")
			} else {
				fmt.Println(err)
			}
			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"`)) {
			fmt.Println("TLS:", statusOK)
			checkCertExpiry(c)
		} else {
			fmt.Println("TLS:", statusNOK)
			fmt.Println("XMPP stream negotiation failed.")
		}
	}
}

func checkCertExpiry(c *tls.Conn) {
	expiry := c.ConnectionState().PeerCertificates[0].NotAfter
	start := c.ConnectionState().PeerCertificates[0].NotBefore
	now := time.Now()
	if now.Before(expiry) && now.After(start) {
		fmt.Println("Certificate:", statusOK)
	} else {
		fmt.Println("Certificate:", statusNOK)
		fmt.Println("Valid from", start, "to", expiry)
	}
}

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
}
