kittemail/main.go

package main

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"crypto/ed25519"
	"crypto/tls"
	"encoding/json"
	"git.ailur.dev/ailur/smtp"
	"github.com/emersion/go-imap/server"
	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"net/http"
	"net/textproto"

	library "git.ailur.dev/ailur/fg-library/v2"
	nucleusLibrary "git.ailur.dev/ailur/fg-nucleus-library"
)

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

// Log a message to the logger service
// messageType:
// 0 - Information
// 1 - Warning
// 2 - Error
// 3 - Guru (exits immediately)
func log(message string, messageType uint64) {
	Information.Outbox <- library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000002"),
		MessageType:  messageType,
		SentAt:       time.Now(),
		Message:      message,
	}
}

// Authenticate Fake for testing
func Authenticate(token string, config OAuthConfig) (uuid.UUID, error) {
	println("called")
	return uuid.MustParse("e59fece6-256f-4799-bb31-321268387d12"), nil
}

// GetUsername Fake for testing
func GetUsername(token string, config 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 []byte, owner uuid.UUID) error {
	response, err := timedOutInterServiceMessage(library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000003"),
		MessageType:  3,
		SentAt:       time.Now(),
		Message: nucleusLibrary.File{
			Name:  name,
			User:  owner,
			Bytes: data,
		},
	}, 3*time.Second)
	if err != nil {
		return err
	}

	switch response.MessageType {
	case 0:
		return nil
	case 1, 2:
		return response.Message.(error)
	default:
		return errors.New("unknown error")
	}
}

func GetFile(name string, owner uuid.UUID) (*os.File, error) {
	response, err := timedOutInterServiceMessage(library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000003"),
		MessageType:  4,
		SentAt:       time.Now(),
		Message: nucleusLibrary.File{
			Name: name,
			User: owner,
		},
	}, 3*time.Second)
	if err != nil {
		return nil, err
	}

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

		return file, nil
	case 1, 2:
		return nil, errors.New(response.Message.(string))
	default:
		return nil, errors.New("unknown error")
	}
}

func DeleteFile(name string, owner uuid.UUID) error {
	response, err := timedOutInterServiceMessage(library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000003"),
		MessageType:  5,
		SentAt:       time.Now(),
		Message: nucleusLibrary.File{
			Name: name,
			User: owner,
		},
	}, 3*time.Second)
	if err != nil {
		return err
	}

	switch response.MessageType {
	case 0:
		return nil
	case 1, 2:
		return response.Message.(error)
	default:
		return errors.New("unknown error")
	}
}

func timedOutInterServiceMessage(message library.InterServiceMessage, timeout time.Duration) (response library.InterServiceMessage, err error) {
	Information.Outbox <- message

	var stopped sync.Once
	stop := make(chan struct{})

	go func() {
		time.Sleep(timeout)
		if response == (library.InterServiceMessage{}) {
			err = errors.New("timed out")
			stopped.Do(func() {
				close(stop)
			})
		}
	}()

	go func() {
		for {
			select {
			case response = <-Information.Inbox:
				stopped.Do(func() {
					close(stop)
				})
				return
			case <-stop:
				return
			}
		}
	}()

	select {
	case <-stop:
		return
	}
}

func getDatabase() (database library.Database, err error) {
	var response library.InterServiceMessage
	response, err = timedOutInterServiceMessage(library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000001"),
		MessageType:  1,
		SentAt:       time.Now(),
		Message:      nil,
	}, 3*time.Second)
	if err != nil {
		return
	}

	switch response.MessageType {
	case 2:
		var ok bool
		database, ok = response.Message.(library.Database)
		if !ok {
			err = errors.New("database not found")
			return
		}

		err = setupDatabase(database)
		if err != nil {
			return
		}
	case 1, 0:
		err = errors.New(response.Message.(string))
	default:
		err = errors.New("unknown error")
	}

	return
}

func getPublicKey() (publicKey ed25519.PublicKey, err error) {
	var response library.InterServiceMessage
	response, err = timedOutInterServiceMessage(library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000004"),
		MessageType:  2,
		SentAt:       time.Now(),
		Message:      nil,
	}, 3*time.Second)
	if err != nil {
		return
	}

	switch response.MessageType {
	case 2:
		var ok bool
		publicKey, ok = response.Message.(ed25519.PublicKey)
		if !ok {
			err = errors.New("publicKey not found")
			return
		}
	default:
		err = errors.New("unknown error")
	}

	return
}

func getOAuthHostName() (hostName string, err error) {
	var response library.InterServiceMessage
	response, err = timedOutInterServiceMessage(library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000004"),
		MessageType:  0,
		SentAt:       time.Now(),
		Message:      nil,
	}, 3*time.Second)
	if err != nil {
		return
	}

	switch response.MessageType {
	case 0:
		var ok bool
		hostName, ok = response.Message.(string)
		if !ok {
			err = errors.New("oauthHostName not found")
			return
		}
	case 1, 2:
		err = response.Message.(error)
	default:
		err = errors.New("unknown error")
	}

	return
}

func registerOAuth(hostName string) (oauthRegistration nucleusLibrary.OAuthResponse, err error) {
	var urlPath string
	urlPath, err = url.JoinPath(hostName, "/oauth")
	if err != nil {
		return
	}

	var response library.InterServiceMessage
	response, err = timedOutInterServiceMessage(library.InterServiceMessage{
		ServiceID:    ServiceInformation.ServiceID,
		ForServiceID: uuid.MustParse("00000000-0000-0000-0000-000000000004"),
		MessageType:  1,
		SentAt:       time.Now(),
		Message: nucleusLibrary.OAuthInformation{
			Name:        "Kittemail",
			RedirectUri: urlPath,
			Scopes:      []string{"openid"},
		},
	}, 3*time.Second)
	if err != nil {
		return
	}

	switch response.MessageType {
	case 0:
		var ok bool
		oauthRegistration, ok = response.Message.(nucleusLibrary.OAuthResponse)
		if !ok {
			err = errors.New("oauthRegistration not found")
		}
	case 1, 2:
		err = errors.New(response.Message.(string))
	default:
		err = errors.New("unknown error")
	}

	return
}

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

	publicKey, err = getPublicKey()
	if err != nil {
		return
	}

	oauthHostName, err = getOAuthHostName()
	if err != nil {
		return
	}

	oauthRegistration, err = registerOAuth(hostName)
	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, flags TEXT NOT NULL, mailbox BLOB NOT NULL, FOREIGN KEY (mailbox) REFERENCES mailboxes(id))")
		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, flags TEXT NOT NULL, mailbox BYTEA NOT NULL, FOREIGN KEY (mailbox) REFERENCES mailboxes(id))")
		if err != nil {
			return err
		}
	}

	return nil
}

type ByteLiteral struct {
	UnderlyingReader io.Reader
	Length           int
}

func (l *ByteLiteral) Read(p []byte) (n int, err error) {
	return l.UnderlyingReader.Read(p)
}

func (l *ByteLiteral) Len() int {
	return l.Length
}

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)),
					openMessages: make(map[*Message]struct{}),
				}

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

				err = mailbox.CreateMessage(nil, time.Now(), &ByteLiteral{
					UnderlyingReader: bytes.NewReader(mail.Data),
					Length:           len(mail.Data),
				})
				if err != nil {
					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()

	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 := &Backend{oauthConfig: 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() {
		imapServer := server.New(imapBackend)
		if !enforceTLS {
			imapServer.AllowInsecureAuth = true
		}
		imapServer.TLSConfig = tlsConfig
		imapServer.Addr = listenerHost + ":143"
		err := imapServer.ListenAndServe()
		if err != nil {
			log("Failed to serve IMAP: "+err.Error(), 3)
		}
	}()

	// Implicit TLS IMAP port
	go func() {
		imapServer := server.New(imapBackend)
		if !enforceTLS {
			imapServer.AllowInsecureAuth = true
		}
		imapServer.TLSConfig = tlsConfig
		imapServer.Addr = listenerHost + ":993"
		err := imapServer.ListenAndServeTLS()
		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)
		}
	}()
}