reactive-mail-server/backend/src/main/java/com/mailserver/imap/ImapWriter.java

package com.mailserver.imap;

import jakarta.mail.internet.AddressException;
import jakarta.mail.internet.ContentDisposition;
import jakarta.mail.internet.InternetAddress;
import jakarta.mail.internet.MimeUtility;
import jakarta.mail.internet.ParameterList;
import jakarta.mail.internet.ParseException;

import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
import java.util.Enumeration;
import java.util.Locale;

/** Builds IMAP response bytes: strings, literals, ENVELOPE, BODYSTRUCTURE (RFC 3501 7.4.2). */
final class ImapWriter {

    private final ByteArrayOutputStream out = new ByteArrayOutputStream();

    ImapWriter text(String text) {
        byte[] bytes = text.getBytes(StandardCharsets.UTF_8);
        out.write(bytes, 0, bytes.length);
        return this;
    }

    ImapWriter sp() {
        out.write(' ');
        return this;
    }

    ImapWriter crlf() {
        return text("\r\n");
    }

    ImapWriter literal(byte[] bytes) {
        text("{" + bytes.length + "}\r\n");
        out.write(bytes, 0, bytes.length);
        return this;
    }

    /** Quoted when it is plain ASCII without CR/LF, a literal otherwise; null → NIL. */
    ImapWriter nstring(String value) {
        if (value == null) {
            return text("NIL");
        }
        return string(value);
    }

    ImapWriter string(String value) {
        boolean quotable = value.length() < 1024;
        for (int i = 0; i < value.length() && quotable; i++) {
            char c = value.charAt(i);
            if (c > 0x7e || c < 0x20) {
                quotable = false;
            }
        }
        if (!quotable) {
            return literal(value.getBytes(StandardCharsets.UTF_8));
        }
        return text("\"" + value.replace("\\", "\\\\").replace("\"", "\\\"") + "\"");
    }

    byte[] bytes() {
        return out.toByteArray();
    }

    // ---------- ENVELOPE ----------

    ImapWriter envelope(MimeTree message) {
        text("(");
        nstring(message.header("Date")).sp();
        nstring(message.header("Subject")).sp();
        String from = message.header("From");
        addresses(from).sp();
        addresses(firstNonNull(message.header("Sender"), from)).sp();
        addresses(firstNonNull(message.header("Reply-To"), from)).sp();
        addresses(message.header("To")).sp();
        addresses(message.header("Cc")).sp();
        addresses(message.header("Bcc")).sp();
        nstring(message.header("In-Reply-To")).sp();
        nstring(message.header("Message-ID"));
        return text(")");
    }

    private ImapWriter addresses(String header) {
        if (header == null || header.isBlank()) {
            return text("NIL");
        }
        InternetAddress[] list;
        try {
            list = InternetAddress.parseHeader(header, false);
        } catch (AddressException e) {
            return text("NIL");
        }
        if (list.length == 0) {
            return text("NIL");
        }
        text("(");
        for (InternetAddress address : list) {
            String email = address.getAddress() == null ? "" : address.getAddress();
            int at = email.lastIndexOf('@');
            text("(");
            nstring(encodedPersonal(address.getPersonal())).sp().text("NIL").sp();
            nstring(at > 0 ? email.substring(0, at) : email).sp();
            nstring(at > 0 ? email.substring(at + 1) : null);
            text(")");
        }
        return text(")");
    }

    private static String encodedPersonal(String personal) {
        if (personal == null) {
            return null;
        }
        try {
            return MimeUtility.encodeWord(personal, "UTF-8", "B");
        } catch (UnsupportedEncodingException e) {
            return personal;
        }
    }

    // ---------- BODYSTRUCTURE / BODY ----------

    ImapWriter bodyStructure(MimeTree part, boolean extensible) {
        text("(");
        if (part.isMultipart() && !part.children.isEmpty()) {
            for (MimeTree child : part.children) {
                bodyStructure(child, extensible);
            }
            sp().string(part.type.getSubType().toUpperCase(Locale.ROOT));
            if (extensible) {
                sp().params(part.type.getParameterList()).sp();
                disposition(part).sp().text("NIL NIL");
            }
            return text(")");
        }
        String primary = part.isMultipart() ? "TEXT" : part.type.getPrimaryType().toUpperCase(Locale.ROOT);
        String sub = part.isMultipart() ? "PLAIN" : part.type.getSubType().toUpperCase(Locale.ROOT);
        string(primary).sp().string(sub).sp();
        params(part.type.getParameterList()).sp();
        nstring(part.header("Content-ID")).sp();
        nstring(part.header("Content-Description")).sp();
        String encoding = part.header("Content-Transfer-Encoding");
        string(encoding == null ? "7BIT" : encoding.toUpperCase(Locale.ROOT)).sp();
        text(String.valueOf(part.bodySize()));
        if ("TEXT".equals(primary)) {
            sp().text(String.valueOf(part.bodyLines()));
        } else if (part.message != null) {
            sp().envelope(part.message).sp().bodyStructure(part.message, extensible).sp().text(String.valueOf(part.bodyLines()));
        }
        if (extensible) {
            sp().nstring(part.header("Content-MD5")).sp();
            disposition(part).sp();
            nstring(part.header("Content-Language")).sp();
            nstring(part.header("Content-Location"));
        }
        return text(")");
    }

    private ImapWriter params(ParameterList params) {
        if (params == null || params.size() == 0) {
            return text("NIL");
        }
        text("(");
        boolean first = true;
        Enumeration<String> names = params.getNames();
        while (names.hasMoreElements()) {
            String name = names.nextElement();
            if (!first) {
                sp();
            }
            first = false;
            string(name.toUpperCase(Locale.ROOT)).sp().string(String.valueOf(params.get(name)));
        }
        return text(")");
    }

    private ImapWriter disposition(MimeTree part) {
        String header = part.header("Content-Disposition");
        if (header == null) {
            return text("NIL");
        }
        try {
            ContentDisposition disposition = new ContentDisposition(header);
            text("(").string(disposition.getDisposition().toUpperCase(Locale.ROOT)).sp();
            params(disposition.getParameterList());
            return text(")");
        } catch (ParseException e) {
            return text("NIL");
        }
    }

    private static String firstNonNull(String a, String b) {
        return a != null ? a : b;
    }

    // ---------- modified UTF-7 (RFC 3501 5.1.3) for mailbox names ----------

    static String encodeMailbox(String name) {
        StringBuilder out = new StringBuilder();
        StringBuilder pending = new StringBuilder();
        for (int i = 0; i < name.length(); i++) {
            char c = name.charAt(i);
            if (c >= 0x20 && c <= 0x7e) {
                flushBase64(out, pending);
                out.append(c == '&' ? "&-" : String.valueOf(c));
            } else {
                pending.append(c);
            }
        }
        flushBase64(out, pending);
        return out.toString();
    }

    private static void flushBase64(StringBuilder out, StringBuilder pending) {
        if (pending.length() == 0) {
            return;
        }
        byte[] utf16 = pending.toString().getBytes(StandardCharsets.UTF_16BE);
        String base64 = java.util.Base64.getEncoder().withoutPadding().encodeToString(utf16).replace('/', ',');
        out.append('&').append(base64).append('-');
        pending.setLength(0);
    }

    static String decodeMailbox(String name) {
        StringBuilder out = new StringBuilder();
        int i = 0;
        while (i < name.length()) {
            char c = name.charAt(i);
            if (c != '&') {
                out.append(c);
                i++;
                continue;
            }
            int dash = name.indexOf('-', i);
            if (dash < 0) {
                out.append(name.substring(i));
                break;
            }
            if (dash == i + 1) {
                out.append('&');
            } else {
                try {
                    byte[] utf16 = java.util.Base64.getDecoder().decode(name.substring(i + 1, dash).replace(',', '/'));
                    out.append(new String(utf16, StandardCharsets.UTF_16BE));
                } catch (IllegalArgumentException e) {
                    out.append(name, i, dash + 1);
                }
            }
            i = dash + 1;
        }
        return out.toString();
    }
}