// mxrouter
// Copyright (c) 2020 yo000 <johan@nosd.in>
//
// TODO : REconnexion ldap lorsque necessaire

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

package main

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

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

const (
	version = "0.9.4"
)

var (
	logstream  *syslog.Writer
	mutex      sync.Mutex
	debug      *bool
	listen     *string
	ldapURL    *string
	ldapBaseDN *string
	ldapUser   *string
	ldapPass   *string
	ldapAttr   *string
	defDomain  *string
)


func ldapSearch(conLdap *ldap.Conn, 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")
		conLdap.Close()
		conLdap, err = connectLdap()
		if err != nil {
			return result, err
		} else {
			attempt = attempt + 1
			return ldapSearch(conLdap, searchReq, attempt)
		}
	}

	return result, err
}


func handleConnection(connClt net.Conn, conLdap *ldap.Conn) {
	var s string
	var mxdom string
	buf := make([]byte, 1024)

	defer connClt.Close()
	readlen, err := connClt.Read(buf)
	if err != nil {
		fmt.Println("Error reading:", err.Error())
	}

	// 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 := strings.Replace(fmt.Sprintf("Error lookup MX: %s", err.Error()), " ", "%20", -1)
			response := fmt.Sprintf("500 %s\n", s)
			connClt.Write([]byte(response))
			return
		}
		// 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 := strings.Replace("No usable MX found", " ", "%20", -1)
			resp := fmt.Sprintf("500 %s\n", s)
			connClt.Write([]byte(resp))
			return
		}

		mxslic := strings.Split(mxs[0].Host, ".")
		// les MXs terminent par '.', on le retire
		if mxs[0].Host[len(mxs[0].Host)-1] == '.' {
			mxdom = strings.Join(mxslic[len(mxslic)-3:len(mxslic)-1], ".")
		} else {
			mxdom = strings.Join(mxslic[len(mxslic)-2:len(mxslic)], ".")
		}

		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(conLdap, searchReq, 0)
		if err != nil {
			logstream.Err(fmt.Sprintln("Error searching into LDAP: ", err))
			s := strings.Replace(err.Error(), " ", "%20", -1)
			response := fmt.Sprintf("500 %s\n", s)
			connClt.Write([]byte(response))
			return
		}

		// 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 := strings.Replace(fmt.Sprintf("Attribute not found: %s", *ldapAttr), " ", "%20", -1)
			response := fmt.Sprintf("500 %s\n", s)
			connClt.Write([]byte(response))
			return
		}
		if len(result.Entries) != 1 {
			if *debug {
				logstream.Debug("Got no result, returning 500")
			}
			s := strings.Replace("No route defined", " ", "%20", -1)
			response := fmt.Sprintf("500 %s\n", s)
			connClt.Write([]byte(response))
			return
		}
		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 = strings.Replace(fmt.Sprintf("FILTER relay:[%s]", result.Entries[0].Attributes[0].Values[0]), " ", "%20", -1)
		} else {
			s = strings.Replace(fmt.Sprintf("FILTER relay:[%s]", fmt.Sprintf("%s.%s", result.Entries[0].Attributes[0].Values[0], *defDomain)), " ", "%20", -1)
		}
		connClt.Write([]byte(fmt.Sprintf("200 %s\n", s)))
	} else {
		if *debug {
			logstream.Debug("Incorrect input format : " + string(buf[:readlen-1]))
		}
		s := strings.Replace("Incorrect input format", " ", "%20", -1)
		response := fmt.Sprintf("500 %s\n", s)
		connClt.Write([]byte(response))
		return
	}
}

func connectLdap() (*ldap.Conn, 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)")
	_ = 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 logstream, e = syslog.New(syslog.LOG_MAIL, "mxrouter"); e != nil {
		log.Fatal(e)
	}
	defer logstream.Close()
	run()
}