#!/usr/bin/env python3
# Copyright (c) Microsoft Corporation.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, see <https://www.gnu.org/licenses/>.

import imaplib
import gssapi

from socket import getfqdn

IMAP_HOST = getfqdn()
SERVICE_NAME = "imap"

imaplib.Debug = 4

def gssapi_authenticate(imap_conn: imaplib.IMAP4, host: str, service_name: str = "imap"):
    """
    Perform SASL GSSAPI authentication against an already-connected
    imaplib.IMAP4 (or IMAP4_SSL) instance.
    """

    # Build the target service principal, e.g. "imap@imap.example.com"
    service = gssapi.Name(f"{service_name}@{host}", gssapi.NameType.hostbased_service)

    # Establish a GSSAPI security context for that service.
    ctx = gssapi.SecurityContext(name=service, usage="initiate")

    # State shared across callback invocations. IMAP SASL exchange is
    # multi-step, so imaplib will call our callback repeatedly with
    # each server challenge until authentication completes.
    state = {"step": 0}

    def sasl_callback(challenge: bytes) -> bytes:
        step = state["step"]
        state["step"] += 1

        if not ctx.complete:
            # Step 1..N: perform the GSSAPI negotiation.
            # On the first call, challenge is typically empty.
            out_token = ctx.step(challenge)
            if out_token is None:
                out_token = b""
            return out_token

        # Context is established: final step handles the security
        # layer negotiation (RFC 2222 / RFC 4752).
        # The server sends an encrypted message indicating supported
        # security layers and max message size; we unwrap it, then
        # respond with our chosen layer (usually "no security layer",
        # i.e. byte 1) and, optionally, our username/authorization id.
        unwrapped = ctx.unwrap(challenge).message

        # unwrapped[0] = bitmask of supported layers
        # unwrapped[1:4] = max buffer size the server can accept
        # We select "no security layer" (0x01) and keep buffer size.
        response = bytes([1]) + unwrapped[1:4]

        # Optionally append the authorization identity (often the
        # username) to the response, per RFC 4752.
        # response += authzid.encode("utf-8")

        wrapped = ctx.wrap(response, False).message
        return wrapped

    # imaplib's authenticate() expects a callable that takes the raw
    # (already base64-decoded) challenge and returns raw bytes to send
    # (imaplib handles the base64 encoding/decoding for you).
    typ, data = imap_conn.authenticate("GSSAPI", sasl_callback)

    if typ != "OK":
        raise imaplib.IMAP4.error(f"GSSAPI authentication failed: {data}")

    return typ, data


def main():
    imap_conn = imaplib.IMAP4(IMAP_HOST)

    try:
        typ, data = gssapi_authenticate(imap_conn, IMAP_HOST, SERVICE_NAME)
        print("Authenticated:", typ, data)

        # Example: list mailboxes after successful auth.
        typ, mailboxes = imap_conn.list()
        print(typ, mailboxes)

        # Example: select INBOX and get message count.
        typ, sel_data = imap_conn.select("INBOX")
        print("Selected INBOX:", typ, sel_data)

    finally:
        try:
            imap_conn.logout()
        except Exception:
            pass


if __name__ == "__main__":
    main()
