#!/usr/bin/env python3
# List folders on an IMAP server.
import argparse
import gssapi # for GSSError
import imaplib
from nullroute.core import Core
from nullroute.sec.sasl import SaslGSSAPI
from nullroute.string.imaputf7 import encode_imap_utf7, decode_imap_utf7
import re
import ssl
import sys
import urllib.parse

def find_default_domain():
    import socket

    fqdn = socket.getfqdn()
    if "." not in fqdn:
        raise RuntimeError("local host does not have a domain name")

    _, domain = fqdn.split(".", 1)
    return domain

def find_default_host(want_starttls=False):
    return discover_srv(find_default_domain(),
                        want_starttls=want_starttls)

def discover_srv(domain, *, want_starttls=False):
    """
    Find IMAP server for specified domain via DNS SRV.

    If 'want_starttls' is given, will only search for the 'imap' service.
    By default will search for 'imaps' first, falling back to 'imap'.
    """

    import dns.resolver
    import dns.rdtypes.util # Workaround for missing import

    if want_starttls:
        Core.debug("querying for service 'imap'")
        starttls = True
        answer = dns.resolver.resolve(f"_imap._tcp.{domain}", "SRV")
    else:
        try:
            Core.debug("querying for service 'imaps'")
            starttls = False
            answer = dns.resolver.resolve(f"_imaps._tcp.{domain}", "SRV")
        except dns.resolver.NXDOMAIN:
            Core.debug("querying for service 'imap'")
            starttls = True
            answer = dns.resolver.resolve(f"_imap._tcp.{domain}", "SRV")

    try:
        hosts = answer.rrset.processing_order()
    except AttributeError:
        # It seems dnspython 2.0 doesn't have this method, bodge something
        # together as we only use the 1st result for now.
        hosts = sorted(answer.rrset, key=lambda r: r.priority)
    Core.trace("SRV hosts: %r", hosts)

    host = hosts[0].target
    port = hosts[0].port
    Core.debug("found server %r on port %r (starttls=%r)", host, port, starttls)

    return str(host).rstrip("."), port, starttls

def imap_parse_list(resp):
    ok, data = resp
    if ok != "OK":
        Core.die("failed to obtain folder list: %s", data[0].decode())
    for line in data:
        # imaplib only has a half-assed parser so each array item looks like this:
        # b'(\\HasNoChildren \\UnMarked) "/" Archive/Work/LITNET'
        # b'() "/" "Archive/Foobar &2D3cTdg93E7YPdxKJwonDNg93EwnCw-"'
        Core.trace("have line: %r", line)
        m = re.match(br'^\((?P<flags>.*?)\) "(?P<delimiter>.*?)" (?P<name>.*)$', line)
        if m:
            flags, delim, name = m.groups()
            # Names with spaces may be in quotes.
            # In any case the name will be imap utf7 encoded.
            if re.match(br'^".+"$', name):
                name = name[1:-1]
            name = decode_imap_utf7(name)
            yield name
        else:
            Core.die("could not parse IMAP response line: %r", line)

parser = argparse.ArgumentParser()
parser.add_argument("-H", "--host",
                    help="IMAP server to access")
parser.add_argument("--port", type=int,
                    help="Use an alternate port")
parser.add_argument("--starttls", action="store_true",
                    help="Negotiate STARTTLS instead of using TLS implicitly")
parser.add_argument("--cleartext", action="store_true",
                    help="Do not use TLS at all")
parser.add_argument("--username",
                    help="Username for basic authentication (IMAP LOGIN)")
parser.add_argument("--password",
                    help="Password for basic authentication (IMAP LOGIN)")
parser.add_argument("-s", "--lsub", action="store_true",
                    help="Show subscribed folders (LSUB) instead of all folders")
parser.add_argument("--create", metavar="FOLDER",
                    help="Create a folder")
parser.add_argument("--rename", metavar="FOLDER:FOLDER",
                    help="Rename a folder")
parser.add_argument("--delete", metavar="FOLDER",
                    help="Delete a folder")
parser.add_argument("--subscribe", metavar="FOLDER",
                    help="Subscribe to a folder")
parser.add_argument("--unsubscribe", metavar="FOLDER",
                    help="Unsubscribe from a folder")
args = parser.parse_args()

if args.host and "://" in args.host:
    url = urllib.parse.urlparse(args.host)
    if url.scheme not in {"imap", "imaps"}:
        Core.die("non-IMAP URL was specified")
    args.host = url.hostname
    args.starttls = (url.scheme == "imap")
    if not args.port:
        args.port = url.port
    if not args.username:
        args.username = urllib.parse.unquote(url.username)
    if not args.password:
        args.password = urllib.parse.unquote(url.password)

if not args.host:
    if args.port:
        Core.die("contradictory options (--port but no --host) given")
    if args.cleartext:
        args.host, args.port, _ = find_default_host(True)
    else:
        args.host, args.port, args.starttls = find_default_host(args.starttls)

if args.cleartext and args.password:
    # Only GSSAPI is supported. (Maybe SCRAM in the future.)
    Core.die("refusing to do plain-password LOGIN on a cleartext connection")

tlsctx = ssl.create_default_context()

try:
    if args.starttls and args.cleartext:
        Core.die("contradictory options (--starttls and --cleartext) given")
    elif args.starttls or args.cleartext:
        Core.debug("connecting to %r using plaintext", args.host)
        clnt = imaplib.IMAP4(args.host,
                             args.port or imaplib.IMAP4_PORT)
        if args.starttls:
            Core.debug("negotiating STARTTLS")
            clnt.starttls(ssl_context=tlsctx)
    else:
        Core.debug("connecting to %r using TLS", args.host)
        clnt = imaplib.IMAP4_SSL(args.host,
                                 args.port or imaplib.IMAP4_SSL_PORT,
                                 ssl_context=tlsctx)
except ssl.SSLCertVerificationError as e:
    Core.die("TLS connection failed: %s", e)

if args.username and args.password:
    Core.debug("authenticating as %r using plaintext", args.username)
    try:
        clnt.login(args.username, args.password)
    except imaplib.IMAP4.error as e:
        Core.die("LOGIN authentication failed: %s", e)
else:
    sasl = SaslGSSAPI(args.host, "imap")
    try:
        Core.debug("authenticating using %s", sasl.mech_name)
        clnt.authenticate(sasl.mech_name, sasl)
    except gssapi.raw.misc.GSSError as e:
        Core.die("GSSAPI authentication failed: %s", e)

if args.create:
    mbox = args.create
    Core.info("creating mailbox %r", mbox)
    ok, data = clnt.create('"%s"' % encode_imap_utf7(mbox).decode())
    if ok != "OK":
        Core.err("failed to create folder %r: %s", mbox, data[0].decode())
    clnt.logout()
elif args.rename:
    oldmbox, newmbox = args.rename.split(":")
    Core.info("renaming mailbox %r to %r", oldmbox, newmbox)
    ok, data = clnt.unsubscribe('"%s"' % encode_imap_utf7(oldmbox).decode())
    if ok != "OK":
        Core.err("failed to unsubscribe from folder %r: %s", oldmbox, data[0].decode())
    ok, data = clnt.rename('"%s"' % encode_imap_utf7(oldmbox).decode(),
                           '"%s"' % encode_imap_utf7(newmbox).decode())
    if ok != "OK":
        Core.err("failed to rename folder %r to %r: %s", oldmbox, newmbox, data[0].decode())
    ok, data = clnt.subscribe('"%s"' % encode_imap_utf7(newmbox).decode())
    if ok != "OK":
        Core.err("failed to subscribe to folder %r: %s", newmbox, data[0].decode())
    clnt.logout()
elif args.delete:
    mbox = args.delete
    Core.info("deleting mailbox %r", mbox)
    ok, data = clnt.unsubscribe('"%s"' % encode_imap_utf7(mbox).decode())
    if ok != "OK":
        Core.err("failed to unsubscribe from folder %r: %s", mbox, data[0].decode())
    ok, data = clnt.delete('"%s"' % encode_imap_utf7(mbox).decode())
    if ok != "OK":
        Core.err("failed to delete folder %r: %s", mbox, data[0].decode())
    clnt.logout()
elif args.subscribe:
    mbox = args.subscribe
    Core.info("subscribing to mailbox %r", mbox)
    ok, data = clnt.subscribe('"%s"' % encode_imap_utf7(mbox).decode())
    if ok != "OK":
        Core.err("failed to subscribe to folder %r: %s", mbox, data[0].decode())
    clnt.logout()
elif args.unsubscribe:
    mbox = args.unsubscribe
    Core.info("unsubscribing from mailbox %r", mbox)
    ok, data = clnt.unsubscribe('"%s"' % encode_imap_utf7(mbox).decode())
    if ok != "OK":
        Core.err("failed to unsubscribe from folder %r: %s", mbox, data[0].decode())
    clnt.logout()
else:
    all_folders = sorted(imap_parse_list(clnt.list()))
    sub_folders = sorted(imap_parse_list(clnt.lsub()))
    clnt.logout()
    if args.lsub:
        if sys.stdout.isatty():
            # Highlight nonexistent subscribed folders in red (Mutt likes to
            # complain about those).
            all_folders = {*all_folders}
            sub_folders = [*sub_folders]
            for item in sub_folders:
                if item in all_folders:
                    print(item)
                else:
                    print("\033[38;5;1m%s\033[m" % item)
        else:
            for item in sub_folders:
                print(item)
    else:
        if sys.stdout.isatty():
            # Dim non-subscribed folders.
            all_folders = [*all_folders]
            sub_folders = {*sub_folders}
            for item in all_folders:
                if item in sub_folders:
                    print(item)
                else:
                    print("\033[2m%s\033[m" % item)
        else:
            for item in all_folders:
                print(item)
Core.fini()
