// mxrouter
// Copyright (c) 2020 yo000 <johan@nosd.in>
//

// checkmx : postfix tcp_table server which returns second-level domain name of the MX of a domain
// ex :
// "get cfnews.net"
// returns : 200 google.com
//
// Use a single LDAP connection protected by a mutex
//
// v0.9.6 : add PID file
// v0.9.8 : Check MX format with regex
//

package main

import (
	"flag"
	"fmt"
	"log"
	"log/syslog"
	"net"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/go-ldap/ldap/v3"
	"github.com/peterbourgon/ff"
	"github.com/tabalt/pidfile"
)

const (
	version = "0.9.8"
	// Match le nom de domaine, soit les 2 dernieres chaines separees par un point.
	//  Peut aussi terminer par un point (Ex: srv.domaine.com. => domaine.com)
	DomainRegexpFormat = `([0-9A-Za-z\-]*)\.([0-9A-Za-z\-]*)(?:\.)?$`
)

var (
	logstream   *syslog.Writer
	conLdap     *ldap.Conn
	mutex       sync.Mutex
	debug       *bool
	listen      *string
	ldapURL     *string
	ldapBaseDN  *string
	ldapUser    *string
	ldapPass    *string
	ldapAttr    *string
	defDomain   *string
	pidFilePath *string
	timeout     *int
)

func ldapSearch(searchReq *ldap.SearchRequest, attempt int) (*ldap.SearchResult, error) {
	mutex.Lock()
	result, err := conLdap.Search(searchReq)
	mutex.Unlock()
	// Let's just manage connection errors here
	if err != nil && strings.HasSuffix(err.Error(), "ldap: connection closed") {
		logstream.Err("LDAP connection closed, retrying")
		mutex.Lock()
		conLdap.Close()
		conLdap, err = connectLdap()
		mutex.Unlock()
		if err != nil {
			return result, err
		} else {
			attempt = attempt + 1
			return ldapSearch(searchReq, attempt)
		}
	}

	return result, err
}

func sendResponse(con net.Conn, respMsg string, respCode int) error {
	response := fmt.Sprintf("%d %s\n", respCode, strings.Replace(respMsg, " ", "%20", -1))
	_, err := con.Write([]byte(response))
	if err != nil {
		return err
	}

	return nil
}

func handleConnection(connClt net.Conn, conLdap *ldap.Conn) {
	var s string
	var mxdom string

	buf := make([]byte, 1024)

	re := regexp.MustCompile(DomainRegexpFormat)

	// Close connection when this function ends
	defer func() {
		logstream.Debug("Closing connection")
		connClt.Close()
	}()

	timeoutDuration := time.Duration(*timeout) * time.Second

	for {
		// Set a deadline for reading. Read operation will fail if no data
		// is received after deadline.
		connClt.SetReadDeadline(time.Now().Add(timeoutDuration))

		readlen, err := connClt.Read(buf)
		if err != nil {
			// Dont notice if client closed connection
			if err.Error() != "EOF" && !strings.HasSuffix(err.Error(), "i/o timeout") {
				logstream.Err(fmt.Sprintln("Error reading connection :", err.Error()))
			}
			return
		}

		// 1) Recupere le MX du domaine de l'adresse mail recue
		if strings.HasPrefix(string(buf[:readlen-1]), "get ") && strings.Contains(string(buf[:readlen-1]), "@") {
			mail := string(buf[4 : readlen-1])
			domain := strings.Split(mail, "@")[1]
			mxs, err := net.LookupMX(domain)
			if err != nil {
				logstream.Err(fmt.Sprintln("Error lookup mx: ", err))
				s := fmt.Sprintf("Error lookup MX: %s", err.Error())
				sendResponse(connClt, s, 500)
				continue
			}
			// 2) Requete LDAP pour voir si le domaine de second niveau du MX possede un routage particulier
			// Protection contre ca : example.org.            72      IN      MX      0 .
			// Considerons qu'il n'y aura jamais de "one letter TLD"
			if len(mxs[0].Host) < 4 {
				logstream.Err(fmt.Sprintln("No usable mx found"))
				s := "No usable MX found"
				sendResponse(connClt, s, 500)
				continue
			}

			// On recuperer le nom de domaine de 1er niveau
			group := re.FindSubmatch([]byte(mxs[0].Host))
			if len(group) < 3 || len(group[2]) == 0 {
				logstream.Err(fmt.Sprintln("MX format incorrect: ", mxs[0].Host))
				s := fmt.Sprintf("MX format incorrect: %s", mxs[0].Host)
				sendResponse(connClt, s, 500)
				continue
			}
			mxdom = string(group[1]) + "." + string(group[2])

			filter := fmt.Sprintf("(dc=%s)", ldap.EscapeFilter(mxdom))
			searchReq := ldap.NewSearchRequest(*ldapBaseDN, ldap.ScopeWholeSubtree, 0, 0, 0,
				false, filter, []string{*ldapAttr}, []ldap.Control{})

			result, err := ldapSearch(searchReq, 0)
			if err != nil {
				logstream.Err(fmt.Sprintln("Error searching into LDAP: ", err))
				s := err.Error()
				sendResponse(connClt, s, 500)
				continue
			}

			if len(result.Entries) != 1 {
				if *debug {
					logstream.Debug("Got no result, returning 500")
				}
				s := "No route defined"
				sendResponse(connClt, s, 500)
				continue
			}

			// L'attribut n'existe pas
			if len(result.Entries[0].Attributes) == 0 {
				logstream.Err(fmt.Sprintf("Error searching into LDAP: Attribute %s not found for entry %s\n", *ldapAttr, result.Entries[0]))
				s := fmt.Sprintf("Attribute not found: %s", *ldapAttr)
				sendResponse(connClt, s, 500)
				continue
			}

			if *debug {
				logstream.Debug("Got result, returning " + result.Entries[0].Attributes[0].Values[0])
			}
			// Check if result is a FQDN, if not append defDomain
			if strings.Contains(string(result.Entries[0].Attributes[0].Values[0]), ".") {
				s = fmt.Sprintf("FILTER relay:[%s]", result.Entries[0].Attributes[0].Values[0])
			} else {
				s = fmt.Sprintf("FILTER relay:[%s]", fmt.Sprintf("%s.%s", result.Entries[0].Attributes[0].Values[0], *defDomain))
			}
			sendResponse(connClt, s, 200)
		} else {
			if *debug {
				logstream.Debug("Incorrect input format : " + string(buf[:readlen-1]))
			}
			s := "Incorrect input format"
			sendResponse(connClt, s, 500)
			continue
		}
	}
}

func connectLdap() (*ldap.Conn, error) {
	var err error

	conLdap, err = ldap.DialURL(*ldapURL)
	if err != nil {
		logstream.Err(fmt.Sprintln("Error dialing LDAP: ", err))
		return conLdap, err
	}

	err = conLdap.Bind(*ldapUser, *ldapPass)
	if err != nil {
		logstream.Err(fmt.Sprintln("Error binding LDAP: ", err))
		return conLdap, err
	}
	return conLdap, err
}

func run() {
	logstream.Info("start")
	defer logstream.Info("exit")
	listener, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatal(fmt.Sprintln("Error listening: ", err))
	}

	conLdap, err := connectLdap()
	if err != nil {
		return
	}

	for {
		connClt, err := listener.Accept()
		if err != nil {
			logstream.Err(fmt.Sprintln("Error accepting: ", err))
		}
		go handleConnection(connClt, conLdap)
	}
	conLdap.Close()
}

func main() {
	var e error

	fs := flag.NewFlagSet("mxrouter", flag.ExitOnError)
	listen = fs.String("listen-addr", "127.0.0.1:8080", "listen address for server (also via LISTEN env var)")
	debug = fs.Bool("debug", false, "log debug information (also via DEBUG env var)")
	ldapURL = fs.String("ldap", "", "LDAP Server URL (also via LDAP env var)")
	ldapBaseDN = fs.String("ldapDN", "", "LDAP base DN (also via LDAPDN env var)")
	ldapUser = fs.String("ldapUser", "", "LDAP user DN (also via LDAPUSER env var)")
	ldapPass = fs.String("ldapPass", "", "LDAP user password (also via LDAPPASS env var)")
	ldapAttr = fs.String("ldapAttr", "relayName", "LDAP attribute containing the relay name to return")
	defDomain = fs.String("domain", "", "Domain to add to relay name if not a FQDN (also via DOMAIN env var)")
	pidFilePath = fs.String("pidfile", "/var/run/mxrouter/mxrouter.pid", "PID File (also via PIDFILE env var)")
	timeout = fs.Int("timeout", 5, "timeout in seconds")
	_ = fs.String("config", "", "config file (optional)")

	// Surcharge de la fonction Usage()
	fs.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "%s version %s\n", os.Args[0], version)
		fmt.Fprintf(flag.CommandLine.Output(), "Usage:\n")
		fs.PrintDefaults()
	}
	ff.Parse(fs, os.Args[1:], ff.WithEnvVarNoPrefix(), ff.WithConfigFileFlag("config"), ff.WithConfigFileParser(ff.PlainParser))

	if len(*ldapURL) == 0 || len(*ldapBaseDN) == 0 || len(*ldapUser) == 0 || len(*ldapPass) == 0 {
		fs.Usage()
		return
	}

	if pid, err := pidfile.Create(*pidFilePath); err != nil {
		log.Fatal(err)
	} else {
		defer pid.Clear()
	}

	if logstream, e = syslog.New(syslog.LOG_MAIL, "mxrouter"); e != nil {
		log.Fatal(e)
	}
	defer logstream.Close()
	run()
}