//
// Copyright (c) 2024 Tim Kuijsten <info+git@netsend.nl>
// Copyright (c) 2019 Gilles Chehade <gilles@poolp.org>
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
// ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
// OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
//

package main

import (
	"bufio"
	"flag"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"blitiri.com.ar/go/spf"
	"go.netsend.nl/ossec"
)

type session struct {
	id string
	tm time.Time

	localSender bool

	ip             net.IP
	hostName       string
	heloName       string
	mailFrom       string
	mailFromCalled int

	spfm      sync.Mutex
	spfResult spf.Result
}

var sessions = make(map[string]*session)

var outputChannel chan string

var reporters = map[string]func(*session, []string){
	"link-connect":    linkConnect,
	"link-disconnect": linkDisconnect,
	"link-identify":   linkIdentify,
	"link-auth":       linkAuth,
	"tx-mail":         txMail,
}

var filters = map[string]func(*session, []string){
	"commit": filterCommit,
}

func produceOutput(msgType, sessionId, token, format string, a ...interface{}) {
	var out string

	out = msgType + "|" + sessionId + "|" + token + "|" + fmt.Sprintf(format, a...)

	outputChannel <- out
}

func proceed(sessid string, token string) {
	produceOutput("filter-result", sessid, token, "proceed")
}

func tempReject(sessid string, token string) {
	produceOutput("filter-result", sessid, token, "reject|451 4.4.3 SPF lookup failed")
}

func permReject(sessid string, token, extendedStatusCode string) {
	produceOutput("filter-result", sessid, token, "reject|550 %s", extendedStatusCode)
}

func linkConnect(s *session, params []string) {
	rdns := params[0]
	fcrdns := params[1]
	src := params[2]

	s.tm = time.Now()

	tmp := strings.Split(src, ":")
	tmp = tmp[0 : len(tmp)-1]
	src = strings.Join(tmp, ":")
	if strings.HasPrefix(src, "[") {
		src = src[1 : len(src)-1]
	}

	s.ip = net.ParseIP(src)
	if s.ip == nil {
		s.localSender = true
		return
	}

	if len(rdns) > 0 && fcrdns == "pass" {
		s.hostName = rdns
	}
}

func linkDisconnect(s *session, params []string) {
	delete(sessions, s.id)
}

func linkIdentify(s *session, params []string) {
	s.heloName = params[1]
}

func linkAuth(s *session, params []string) {
	if params[1] == "pass" {
		s.localSender = true
	}
}

func txMail(s *session, params []string) {
	s.mailFromCalled++
	if s.mailFromCalled > 1 {
		if s.mailFrom == params[2] {
			return
		}
	}

	s.mailFrom = params[2]

	s.spfm.Lock()
	go func() {
		s.spfResult, _ = spf.CheckHostWithSender(s.ip, s.heloName, s.mailFrom)
		s.spfm.Unlock()
	}()
}

// Note: depends on filter protocol for thread safety
func (s *session) debug(format string, a ...any) {
	msg := fmt.Sprintf(format, a...)
	fmt.Fprintf(os.Stderr, "%s %s; %s (%s [%s]) %s\n", s.id, msg, s.heloName,
		s.hostName, s.ip, s.mailFrom)
}

func applySpf(s *session, token string) {
	switch s.spfResult {
	case spf.None:
		proceed(s.id, token)
	case spf.Neutral:
		proceed(s.id, token)
	case spf.Pass:
		proceed(s.id, token)
	case spf.Fail:
		s.debug("fail: client not authorized")
		permReject(s.id, token, "5.7.1 SPF fail: client not authorized")
	case spf.SoftFail:
		proceed(s.id, token)
	case spf.TempError:
		s.debug("temperror")
		tempReject(s.id, token)
	case spf.PermError:
		s.debug("permerror")
		permReject(s.id, token, "5.5.2 SPF permerror: invalid SPF record in DNS")
	default:
		panic(s.spfResult)
	}
}

func filterCommit(s *session, params []string) {
	token := params[0]

	if s.localSender {
		s.debug("local session whitelisted")
		proceed(s.id, token)
		return
	}

	if s.spfm.TryLock() {
		// great, spf is done
		applySpf(s, token)
		s.spfm.Unlock()
		return
	}

	// spf still not done, get off main thread
	go func() {
		s.spfm.Lock()
		applySpf(s, token)
		s.spfm.Unlock()
	}()
}

func filterInit() {
	for k := range reporters {
		fmt.Printf("register|report|smtp-in|%s\n", k)
	}
	for k := range filters {
		fmt.Printf("register|filter|smtp-in|%s\n", k)
	}
	fmt.Println("register|ready")
}

func trigger(actions map[string]func(*session, []string), atoms []string) {
	if atoms[4] == "link-connect" {
		// special case to simplify subsequent code
		s := session{}
		s.id = atoms[5]
		sessions[s.id] = &s
	}

	s := sessions[atoms[5]]
	actions[atoms[4]](s, atoms[6:])
}

func skipConfig(scanner *bufio.Scanner) {
	for {
		if !scanner.Scan() {
			os.Exit(0)
		}
		line := scanner.Text()
		if line == "config|ready" {
			return
		}
	}
}

func printusage() {
	fmt.Fprintf(os.Stderr, "usage: %s\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	flag.Usage = printusage
	flag.Parse()

	if flag.NArg() > 0 {
		flag.Usage()
		os.Exit(1)
	}

	err := ossec.PledgePromises("stdio inet dns")
	if err != nil {
		fmt.Fprintf(os.Stderr, "pledge failed: %v\n", err)
		os.Exit(1)
	}

	scanner := bufio.NewScanner(os.Stdin)

	skipConfig(scanner)

	filterInit()

	outputChannel = make(chan string)
	go func() {
		for line := range outputChannel {
			fmt.Println(line)
		}
	}()

	for {
		if !scanner.Scan() {
			os.Exit(0)
		}

		atoms := strings.Split(scanner.Text(), "|")
		if atoms[1] < "0.7" {
			fmt.Fprintf(os.Stderr, "require version >= 0.7, have %s\n", atoms[1])
			os.Exit(1)
		}

		switch atoms[0] {
		case "report":
			trigger(reporters, atoms)
		case "filter":
			trigger(filters, atoms)
		default:
			os.Exit(1)
		}
	}
}