kittemail/main.go

package main

import (
	library "git.ailur.dev/ailur/fg-library/v3"
	nucleusLibrary "git.ailur.dev/ailur/fg-nucleus-library"
	"github.com/emersion/go-imap/v2"
	"github.com/emersion/go-imap/v2/imapserver"
	"github.com/golang-jwt/jwt/v5"

	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"time"

	"crypto/ed25519"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/textproto"

	"git.ailur.dev/ailur/smtp"
	"github.com/google/uuid"
)

var ServiceInformation = library.Service{
	Name: "kittemail",
	Permissions: library.Permissions{
		Authenticate:              true,
		Router:                    false,
		Database:                  true,
		BlobStorage:               true,
		InterServiceCommunication: true,
		Resources:                 false,
	},
	ServiceID: uuid.MustParse("068b0c04-d8c8-4738-90fa-d3827f5abf68"),
}

var (
	loggerService  = uuid.MustParse("00000000-0000-0000-0000-000000000002")
	storageService = uuid.MustParse("00000000-0000-0000-0000-000000000003")
)

// Log a message to the logger service
// messageType:
// 0 - Information
// 1 - Warning
// 2 - Error
// 3 - Guru (exits immediately)
func log(message string, messageType library.MessageCode) {
	// Log the message to the logger service
	Information.SendISMessage(loggerService, messageType, message)
}

// Authenticate Fake for testing
func Authenticate(_ string, _ OAuthConfig) (uuid.UUID, error) {
	println("called")
	return uuid.MustParse("ef585ba0-0ac6-48f8-8ebb-2b6cd97d6a21"), nil
}

// GetUsername Fake for testing
func GetUsername(_ string, _ OAuthConfig) (string, error) {
	return "arzumify", nil
}

func RAuthenticate(token string, config OAuthConfig) (uuid.UUID, error) {
	parsedToken, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
		return config.PublicKey, nil
	})
	if err != nil {
		return uuid.Nil, err
	}

	if !parsedToken.Valid {
		return uuid.Nil, errors.New("invalid token")
	}

	claims, ok := parsedToken.Claims.(jwt.MapClaims)
	if !ok {
		return uuid.Nil, errors.New("invalid token")
	}

	date, err := claims.GetExpirationTime()
	if err != nil || date.Before(time.Now()) || claims["sub"] == nil || claims["isOpenID"] == nil || claims["isOpenID"].(bool) {
		return uuid.Nil, errors.New("invalid token")
	}

	return uuid.MustParse(claims["sub"].(string)), nil
}

func RGetUsername(token string, config OAuthConfig) (string, error) {
	var responseData struct {
		Username string `json:"username"`
	}

	request, err := http.NewRequest("GET", config.HostName+"/api/oauth/userinfo", nil)
	if err != nil {
		return "", err
	}

	request.Header.Set("Authorization", "Bearer "+token)

	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return "", err
	}

	if response.StatusCode != 200 || response.Body == nil || response.Body == http.NoBody {
		return "", errors.New("invalid response")
	}

	err = json.NewDecoder(response.Body).Decode(&responseData)
	if err != nil {
		return "", err
	}

	return responseData.Username, nil
}

func StoreFile(name string, data *io.LimitedReader, owner uuid.UUID) error {
	_, err := Information.SendAndAwaitISMessage(storageService, 3, nucleusLibrary.File{
		Name:   name,
		User:   owner,
		Reader: data,
	}, 3*time.Second)
	if err != nil {
		return err
	}

	return nil
}

func GetFile(name string, owner uuid.UUID) (*os.File, error) {
	response, err := Information.SendAndAwaitISMessage(storageService, 4, nucleusLibrary.File{
		Name: name,
		User: owner,
	}, 3*time.Second)
	if err != nil {
		return nil, err
	}

	file, ok := response.Message.(*os.File)
	if !ok {
		return nil, errors.New("invalid response")
	}

	return file, nil
}

func DeleteFile(name string, owner uuid.UUID) error {
	_, err := Information.SendAndAwaitISMessage(storageService, 5, nucleusLibrary.File{
		Name: name,
		User: owner,
	}, 3*time.Second)
	if err != nil {
		return err
	}

	return nil
}

func beginInitialisation(hostName string) (database library.Database, publicKey ed25519.PublicKey, oauthHostName string, oauthRegistration nucleusLibrary.OAuthResponse, err error) {
	database, err = Information.GetDatabase()
	if err != nil {
		return
	}

	// Initialize the database
	err = setupDatabase(database)
	if err != nil {
		return
	}

	// Initialize the OAuth
	oauthRegistration, publicKey, oauthHostName, err = nucleusLibrary.InitializeOAuth(nucleusLibrary.OAuthInformation{
		Name:        "kittemail",
		RedirectUri: hostName + "/oauth",
		Scopes:      []string{"openid"},
	}, Information)
	if err != nil {
		return
	}

	return
}

func setupDatabase(database library.Database) error {
	if database.DBType == library.Sqlite {
		_, err := database.DB.Exec("CREATE TABLE IF NOT EXISTS emails (id BLOB NOT NULL PRIMARY KEY, prefix TEXT NOT NULL, suffix TEXT NOT NULL, creator BLOB NOT NULL, UNIQUE(prefix, suffix))")
		if err != nil {
			return err
		}

		_, err = database.DB.Exec("CREATE TABLE IF NOT EXISTS mailboxes (id BLOB NOT NULL PRIMARY KEY, owner BLOB NOT NULL, mailbox TEXT NOT NULL, subscribed BOOLEAN NOT NULL DEFAULT FALSE, UNIQUE(mailbox, owner))")
		if err != nil {
			return err
		}

		_, err = database.DB.Exec("CREATE TABLE IF NOT EXISTS messages (id BLOB NOT NULL PRIMARY KEY, owner BLOB NOT NULL, uid INTEGER NOT NULL, created INTEGER NOT NULL, bodySize INTEGER NOT NULL, mailbox BLOB NOT NULL, FOREIGN KEY (mailbox) REFERENCES mailboxes(id))")
		if err != nil {
			return err
		}

		_, err = database.DB.Exec("CREATE TABLE IF NOT EXISTS flags (id BLOB NOT NULL, flag TEXT NOT NULL, FOREIGN KEY (id) REFERENCES messages(id), UNIQUE(id, flag))")
		if err != nil {
			return err
		}
	} else {
		_, err := database.DB.Exec("CREATE TABLE IF NOT EXISTS emails (id BYTEA NOT NULL PRIMARY KEY, prefix TEXT NOT NULL, suffix TEXT NOT NULL, creator BYTEA NOT NULL, UNIQUE(prefix, suffix))")
		if err != nil {
			return err
		}

		_, err = database.DB.Exec("CREATE TABLE IF NOT EXISTS mailboxes (id BYTEA NOT NULL PRIMARY KEY, owner BLOB NOT NULL, mailbox TEXT NOT NULL, subscribed BOOLEAN NOT NULL DEFAULT FALSE, UNIQUE(mailbox, owner))")
		if err != nil {
			return err
		}

		_, err = database.DB.Exec("CREATE TABLE IF NOT EXISTS messages (id BYTEA NOT NULL PRIMARY KEY, owner BYTEA NOT NULL, uid BIGINT NOT NULL, created INTEGER NOT NULL, bodySize BIGINT NOT NULL, mailbox BYTEA NOT NULL, FOREIGN KEY (mailbox) REFERENCES mailboxes(id))")
		if err != nil {
			return err
		}

		_, err = database.DB.Exec("CREATE TABLE IF NOT EXISTS flags (id BYTEA NOT NULL, flag TEXT NOT NULL, FOREIGN KEY (id) REFERENCES messages(id), UNIQUE(id, flag))")
		if err != nil {
			return err
		}
	}

	return nil
}

var (
	Information         *library.ServiceInitializationInformation
	Database            library.Database
	SMTPDatabaseBackend = smtp.DatabaseBackend{
		CheckUser: func(address *smtp.Address) (bool, error) {
			var prefix, suffix string
			err := Database.DB.QueryRow("SELECT prefix, suffix FROM emails WHERE prefix = $1 AND suffix = $2", address.Name, address.Address).Scan(&prefix, &suffix)
			if err != nil {
				return false, errors.New("503 5.5.1 User not found")
			}

			if prefix == address.Name && suffix == address.Address {
				return true, nil
			} else {
				return false, nil
			}
		},
		WriteMail: func(mail *smtp.Mail) error {
			for _, recipient := range mail.To {
				var idRaw []byte
				err := Database.DB.QueryRow("SELECT creator FROM emails WHERE prefix = $1 AND suffix = $2", recipient.Name, recipient.Address).Scan(&idRaw)
				if err != nil {
					return errors.New("503 5.5.1 User not found")
				}

				user := &User{
					sub: uuid.Must(uuid.FromBytes(idRaw)),
				}

				mailbox, err := user.mailbox("INBOX")
				if err != nil {
					return errors.New("503 5.5.1 User not found")
				}

				view := mailbox.NewView()

				_, err = view.appendBuffer(bytes.NewReader(mail.Data), int64(len(mail.Data)), &imap.AppendOptions{
					Time: time.Now(),
				})
				if err != nil {
					err = view.Close()
					if err != nil {
						log("Failed to close view: "+err.Error()+", resource leaks may occur", 1)
					}
					return errors.New("421 4.7.0 Error storing message")
				}

				return nil
			}

			return nil
		},
	}
)

func NewSMTPAuthenticationBackend(OAuthRegistration OAuthConfig) smtp.AuthenticationBackend {
	return smtp.AuthenticationBackend{
		SupportedMechanisms: []string{"PLAIN", "XOAUTH2", "OAUTHBEARER"},
		Authenticate: func(initial string, conn *textproto.Conn) (checkAddr smtp.CheckAddress, finalErr error) {
			initialResponse := strings.SplitN(initial, " ", 2)

			var username, token string
			switch initialResponse[0] {
			case "PLAIN":
				credentials, err := base64.StdEncoding.DecodeString(initialResponse[1])
				if err != nil {
					return nil, errors.New("421 4.7.0 Malformed credentials")
				}

				credentialSlice := bytes.SplitN(bytes.TrimPrefix(credentials, []byte{0x00}), []byte{0x00}, 2)
				username = string(credentialSlice[0])
				token = string(credentialSlice[1])
			case "OAUTHBEARER", "XOAUTH2":
				credentials, err := base64.StdEncoding.DecodeString(initialResponse[1])
				if err != nil {
					return nil, errors.New("421 4.7.0 Malformed credentials")
				}

				credentialSlice := bytes.SplitN(bytes.TrimSuffix(bytes.TrimPrefix(credentials, []byte("user=")), []byte{0x01, 0x01}), []byte{0x01}, 2)
				username = string(credentialSlice[0])
				token = string(credentialSlice[1])
			default:
				return nil, errors.New("503 5.5.1 Invalid authentication method: " + initialResponse[0])
			}

			fmt.Println("Username: " + username)
			fmt.Println("Token: " + token)

			sub, err := Authenticate(token, OAuthRegistration)
			if err != nil {
				return nil, errors.New("421 4.7.0 Invalid credentials")
			}

			usernameCheck, err := GetUsername(token, OAuthRegistration)
			if err != nil {
				return nil, errors.New("421 4.7.0 Invalid credentials")
			}

			if username != usernameCheck {
				return nil, errors.New("421 4.7.0 Username does not match")
			}

			return func(address *smtp.Address) (bool, error) {
				rows, err := Database.DB.Query("SELECT prefix, suffix FROM emails WHERE creator = $1", sub[:])
				if err != nil {
					return false, err
				}

				defer func() {
					err := rows.Close()
					if err != nil {
						log("Failed to close rows: "+err.Error()+", resource leaks may occur", 1)
					}
				}()

				for rows.Next() {
					var prefix, suffix string
					err = rows.Scan(&prefix, &suffix)
					if err != nil {
						return false, err
					}

					if address.Name == prefix && address.Address == suffix {
						return true, nil
					}
				}

				return false, nil
			}, nil
		},
	}
}

func parseConfig() (hostName string, listenerHost string, ownedDomains []string, enforceTLS bool, enableTLS bool, certificatePath string, keyPath string, err error) {
	var ok bool
	hostName, ok = Information.Configuration["hostName"].(string)
	if !ok {
		err = errors.New("hostName not found")
		return
	}

	listenerHost, ok = Information.Configuration["listenerHost"].(string)
	if !ok {
		err = errors.New("listenerHost not found")
		return
	}

	ownedDomains, ok = Information.Configuration["ownedDomains"].([]string)
	if !ok {
		err = errors.New("ownedDomains not found")
		return
	}

	enforceTLS, ok = Information.Configuration["enforceTLS"].(bool)
	if !ok {
		err = errors.New("enforceTLS not found")
		return
	}

	enableTLS, ok = Information.Configuration["enableTLS"].(bool)
	if !ok {
		err = errors.New("enableTLS not found")
		return
	}

	if enforceTLS && !enableTLS {
		err = errors.New("cannot enforce TLS without enabling it")
		return
	}

	if enableTLS {
		tlsConfiguration, ok := Information.Configuration["tlsConfiguration"].(map[string]interface{})
		if !ok {
			err = errors.New("tlsConfiguration not found")
			return
		}

		certificatePath, ok = tlsConfiguration["certificatePath"].(string)
		if !ok {
			err = errors.New("certificatePath not found")
			return
		}

		keyPath, ok = tlsConfiguration["keyPath"].(string)
		if !ok {
			err = errors.New("keyPath not found")
			return
		}
	}

	return
}

func getTLSConfig(enableTLS bool, certificatePath string, keyPath string) (tlsConfig *tls.Config, err error) {
	if enableTLS {
		certificate, err := tls.LoadX509KeyPair(certificatePath, keyPath)
		if err != nil {
			return nil, err
		}

		tlsConfig = &tls.Config{
			Certificates: []tls.Certificate{certificate},
		}
	} else {
		tlsConfig = &tls.Config{
			GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
				return nil, errors.New("TLS disabled in configuration")
			},
		}
	}

	return
}

func Main(information *library.ServiceInitializationInformation) {
	Information = information
	hostName, listenerHost, ownedDomains, enforceTLS, enableTLS, certificatePath, keyPath, err := parseConfig()

	go information.StartISProcessor()

	var oauthConfig OAuthConfig
	//	var oauthRegistration nucleusLibrary.OAuthResponse
	Database, oauthConfig.PublicKey, oauthConfig.HostName, _, err = beginInitialisation(hostName)
	if err != nil {
		log("Failed to initialise: "+err.Error(), 3)
		return
	}

	imapBackend := New(oauthConfig)

	tlsConfig, err := getTLSConfig(enableTLS, certificatePath, keyPath)
	if err != nil {
		log("Failed to get TLS config: "+err.Error(), 3)
		return
	}

	// Plaintext IMAP port
	go func() {
		imapOptions := &imapserver.Options{
			NewSession: func(conn *imapserver.Conn) (imapserver.Session, *imapserver.GreetingData, error) {
				return imapBackend.NewSession(), nil, nil
			},
			Caps: imap.CapSet{
				imap.CapIMAP4rev1: {},
				imap.CapIMAP4rev2: {},
				imap.CapBinary:    {},
			},
			TLSConfig: tlsConfig,
			// TODO: Remove for production
			DebugWriter: os.Stdout,
		}
		if !enforceTLS {
			imapOptions.InsecureAuth = true
		}
		listener, err := net.Listen("tcp", listenerHost+":143")
		if err != nil {
			log("Failed to listen on port 143: "+err.Error(), 3)
			return
		}
		err = imapserver.New(imapOptions).Serve(listener)
		if err != nil {
			log("Failed to serve IMAP: "+err.Error(), 3)
		}
	}()

	// Implicit TLS IMAP port
	go func() {
		imapOptions := &imapserver.Options{
			NewSession: func(conn *imapserver.Conn) (imapserver.Session, *imapserver.GreetingData, error) {
				return imapBackend.NewSession(), nil, nil
			},
			Caps: imap.CapSet{
				imap.CapIMAP4rev1: {},
				imap.CapIMAP4rev2: {},
			},
			TLSConfig: tlsConfig,
			// TODO: Remove for production
			DebugWriter: os.Stdout,
		}
		if !enforceTLS {
			imapOptions.InsecureAuth = true
		}
		listener, err := tls.Listen("tcp", listenerHost+":993", tlsConfig)
		if err != nil {
			log("Failed to listen on port 993: "+err.Error(), 3)
			return
		}
		err = imapserver.New(imapOptions).Serve(listener)
		if err != nil {
			log("Failed to serve IMAP: "+err.Error(), 3)
		}
	}()

	var smtpTLSConfig *tls.Config
	if enableTLS {
		smtpTLSConfig = tlsConfig
	}

	// SMTP Proxy port
	go func() {
		smtpListener, err := net.Listen("tcp", ":25")
		if err != nil {
			log("Failed to listen on port 25: "+err.Error(), 3)
			return
		}

		smtpBackend := smtp.NewReceiver(smtpListener, hostName, ownedDomains, enforceTLS, SMTPDatabaseBackend, NewSMTPAuthenticationBackend(oauthConfig), smtpTLSConfig)
		err = smtpBackend.Serve()
		if err != nil {
			log("Failed to serve SMTP: "+err.Error(), 3)
		}
	}()

	// Plaintext submission port
	go func() {
		smtpListener, err := net.Listen("tcp", ":587")
		if err != nil {
			log("Failed to listen on port 587: "+err.Error(), 3)
			return
		}
		smtpBackend := smtp.NewReceiver(smtpListener, hostName, ownedDomains, enforceTLS, SMTPDatabaseBackend, NewSMTPAuthenticationBackend(oauthConfig), smtpTLSConfig)
		err = smtpBackend.Serve()
		if err != nil {
			log("Failed to serve SMTP submission: "+err.Error(), 3)
		}
	}()

	// Implicit TLS Submission port
	go func() {
		smtpListener, err := net.Listen("tcp", ":465")
		if err != nil {
			log("Failed to listen on port 587: "+err.Error(), 3)
			return
		}
		smtpBackend := smtp.NewReceiver(tls.NewListener(smtpListener, tlsConfig), hostName, ownedDomains, enforceTLS, SMTPDatabaseBackend, NewSMTPAuthenticationBackend(oauthConfig), smtpTLSConfig)
		err = smtpBackend.Serve()
		if err != nil {
			log("Failed to serve SMTP submission: "+err.Error(), 3)
		}
	}()
}