#!/usr/bin/python3
# SPDX-License-Identifier: MPL-2.0
# SPDX-FileCopyrightText: Copyright 2024 Siemens AG

# pylint: disable=missing-docstring,invalid-name

# Renable invalid-name check, it should only cover the module name
# pylint: enable=invalid-name

import argparse
import json
import struct
import sys
import uuid
from enum import Enum

from threading import RLock, Thread
from xml.etree import ElementTree as ET

from gi.repository import GLib
from pydbus import SessionBus
from pydbus.proxy import CompositeInterface

# version is replaced on installation
LINUX_ENTRA_SSO_VERSION = "1.11.0"

# the ssoUrl is a mandatory parameter when requesting a PRT SSO
# Cookie, but the correct value is not checked as of 30.05.2024
# by the authorization backend. By that, a static (fallback)
# value can be used, if no real value is provided.
SSO_URL_DEFAULT = "https://login.microsoftonline.com/"
EDGE_BROWSER_CLIENT_ID = "d7b530a4-7680-4c23-a8bf-c52c121d2e87"
# dbus start service reply codes
START_REPLY_SUCCESS = 1
START_REPLY_ALREADY_RUNNING = 2

# stripped down version of the broker dbus interface,
# as brokers > 2.0.1 do not implement introspection
BROKER_DBUS_SPEC = r"""<!DOCTYPE node PUBLIC
"-//freedesktop//DTD D-Bus Object Introspection 1.0//EN"
"http://www.freedesktop.org/standards/dbus/1.0/introspect.dtd">
<node name="/com/microsoft/identity/broker1">
 <interface name="com.microsoft.identity.Broker1">
  <method name="acquireTokenSilently" >
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="out"/>
  </method>
  <method name="getAccounts" >
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="out"/>
  </method>
  <method name="acquirePrtSsoCookie" >
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="out"/>
  </method>
  <method name="getLinuxBrokerVersion" >
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="in"/>
   <arg type="s" direction="out"/>
  </method>
 </interface>
</node>
"""


class AuthorizationType(Enum):
    CACHED_REFRESH_TOKEN = (1,)
    PRT_SSO_COOKIE = (8,)


class NativeMessaging:
    @staticmethod
    def get_message():
        """
        Read a message from stdin and decode it.
        """
        raw_length = sys.stdin.buffer.read(4)
        if not raw_length:
            sys.exit(0)
        message_length = struct.unpack("@I", raw_length)[0]
        message = sys.stdin.buffer.read(message_length).decode("utf-8")
        return json.loads(message)

    @staticmethod
    def encode_message(message_content):
        """
        Encode a message for transmission, given its content
        """
        encoded_content = json.dumps(message_content, separators=(",", ":")).encode(
            "utf-8"
        )
        encoded_length = struct.pack("@I", len(encoded_content))
        return {"length": encoded_length, "content": encoded_content}

    @staticmethod
    def send_message(encoded_message):
        """
        Send an encoded message to stdout
        """
        sys.stdout.buffer.write(encoded_message["length"])
        sys.stdout.buffer.write(encoded_message["content"])
        sys.stdout.buffer.flush()


class SsoMib:
    BROKER_NAME = "com.microsoft.identity.broker1"
    BROKER_PATH = "/com/microsoft/identity/broker1"
    GRAPH_SCOPES = ["https://graph.microsoft.com/.default"]

    def __init__(self, daemon=False):
        self._bus = SessionBus()
        self.broker = None
        self.session_id = uuid.uuid4()
        self._state_changed_cb = None
        if daemon:
            self._introspect_broker()
            self._monitor_bus()

    def check_broker_online(self):
        return bool(self._bus.get(".DBus").NameHasOwner(self.BROKER_NAME))

    def _introspect_broker(self):
        introspection = ET.fromstring(BROKER_DBUS_SPEC)
        self.broker = CompositeInterface(introspection)(
            self._bus, self.BROKER_NAME, self.BROKER_PATH
        )

    def _monitor_bus(self):
        self._bus.subscribe(
            sender="org.freedesktop.DBus",
            object="/org/freedesktop/DBus",
            signal="NameOwnerChanged",
            arg0=self.BROKER_NAME,
            signal_fired=self._broker_state_changed,
        )

    def _broker_state_changed(
        self, sender, object, iface, signal, params
    ):  # pylint: disable=redefined-builtin,too-many-arguments
        _ = (sender, object, iface, signal)
        # params = (name, old_owner, new_owner)
        new_owner = params[2]
        # the broker introspection is static, so the proxy object never
        # needs to be (re)created or dropped here; just report the state.
        self._report_state_change(bool(new_owner))

    def _report_state_change(self, online):
        if self._state_changed_cb:
            self._state_changed_cb(online)

    def on_broker_state_changed(self, callback):
        """
        Register a callback to be called when the broker state changes.
        The callback should accept a single boolean argument, indicating
        if the broker is online or not.
        """
        self._state_changed_cb = callback

    @staticmethod
    def _get_auth_parameters(account, scopes, sso_url=None):
        params = {
            "account": account,
            "additionalQueryParametersForAuthorization": {},
            "authority": "https://login.microsoftonline.com/common",
            "authorizationType": (
                AuthorizationType.PRT_SSO_COOKIE.value[0]
                if sso_url
                else AuthorizationType.CACHED_REFRESH_TOKEN.value[0]
            ),
            "clientId": EDGE_BROWSER_CLIENT_ID,
            "redirectUri": "https://login.microsoftonline.com"
            "/common/oauth2/nativeclient",
            "requestedScopes": scopes,
            "username": account["username"],
            "uxContextHandle": -1,
        }
        if sso_url:
            params["ssoUrl"] = sso_url
        return params

    def get_accounts(self):
        self._introspect_broker()
        context = {
            "clientId": EDGE_BROWSER_CLIENT_ID,
            "redirectUri": str(self.session_id),
        }
        # pylint: disable=maybe-no-member
        resp = self.broker.getAccounts("0.0", str(self.session_id), json.dumps(context))
        return json.loads(resp)

    def acquire_prt_sso_cookie(
        self, account, sso_url, scopes=GRAPH_SCOPES
    ):  # pylint: disable=dangerous-default-value
        self._introspect_broker()
        request = {
            "account": account,
            "authParameters": SsoMib._get_auth_parameters(account, scopes, sso_url),
            "mamEnrollment": False,
            "ssoUrl": sso_url,
        }
        # pylint: disable=maybe-no-member
        token = json.loads(
            self.broker.acquirePrtSsoCookie(
                "0.0", str(self.session_id), json.dumps(request)
            )
        )
        return token

    def acquire_token_silently(
        self, account, scopes=GRAPH_SCOPES
    ):  # pylint: disable=dangerous-default-value
        self._introspect_broker()
        request = {
            "authParameters": SsoMib._get_auth_parameters(account, scopes),
        }
        # pylint: disable=maybe-no-member
        token = json.loads(
            self.broker.acquireTokenSilently(
                "0.0", str(self.session_id), json.dumps(request)
            )
        )
        return token

    def get_broker_version(self):
        self._introspect_broker()
        params = json.dumps({"msalCppVersion": LINUX_ENTRA_SSO_VERSION})
        # pylint: disable=maybe-no-member
        resp = json.loads(
            self.broker.getLinuxBrokerVersion("0.0", str(self.session_id), params)
        )
        resp["native"] = LINUX_ENTRA_SSO_VERSION
        return resp


def run_as_native_messaging():
    iomutex = RLock()

    def respond(command, message):
        NativeMessaging.send_message(
            NativeMessaging.encode_message({"command": command, "message": message})
        )

    def notify_state_change(online):
        with iomutex:
            respond("brokerStateChanged", "online" if online else "offline")

    def handle_command(cmd, received_message):
        if cmd == "acquirePrtSsoCookie":
            account = received_message["account"]
            sso_url = received_message["ssoUrl"] or SSO_URL_DEFAULT
            token = ssomib.acquire_prt_sso_cookie(account, sso_url)
            respond(cmd, token)
        elif cmd == "acquireTokenSilently":
            account = received_message["account"]
            scopes = received_message.get("scopes") or ssomib.GRAPH_SCOPES
            token = ssomib.acquire_token_silently(account, scopes)
            respond(cmd, token)
        elif cmd == "getAccounts":
            respond(cmd, ssomib.get_accounts())
        elif cmd == "getVersion":
            respond(cmd, ssomib.get_broker_version())

    def run_dbus_monitor():
        # inform other side about the current broker state
        notify_state_change(ssomib.check_broker_online())
        loop = GLib.MainLoop()
        loop.run()

    print("Running as native messaging instance.", file=sys.stderr)
    print("For interactive mode, start with --interactive", file=sys.stderr)

    ssomib = SsoMib(daemon=True)
    ssomib.on_broker_state_changed(notify_state_change)
    # daemon=True is critical: without it, the GLib main loop thread
    # keeps the process alive as an orphan after the browser closes stdin.
    monitor = Thread(target=run_dbus_monitor, daemon=True)
    monitor.start()
    while True:
        received_message = NativeMessaging.get_message()
        with iomutex:
            cmd = received_message["command"]
            try:
                handle_command(cmd, received_message)
            except Exception as exp:  # pylint: disable=broad-except
                err = {"error": f"Failure during request processing: {str(exp)}"}
                respond(cmd, err)


def run_interactive():
    def _get_account(accounts, idx):
        try:
            return accounts["accounts"][idx]
        except IndexError:
            json.dump(
                {"error": f"invalid account index {idx}"},
                indent=2,
                fp=sys.stdout,
            )
            print()
            sys.exit(1)

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        "--interactive",
        action="store_true",
        help="run in interactive mode",
    )
    parser.add_argument(
        "-a",
        "--account",
        type=int,
        default=0,
        help="account index to use for operations",
    )
    parser.add_argument(
        "-s",
        "--ssoUrl",
        default=SSO_URL_DEFAULT,
        help="ssoUrl part of SSO PRT cookie request",
    )
    parser.add_argument(
        "command",
        choices=[
            "getAccounts",
            "getVersion",
            "acquirePrtSsoCookie",
            "acquireTokenSilently",
            "monitor",
        ],
    )
    args = parser.parse_args()

    monitor_mode = args.command == "monitor"
    ssomib = SsoMib(daemon=monitor_mode)
    if monitor_mode:
        print("Monitoring D-Bus for broker availability.")
        ssomib.on_broker_state_changed(
            lambda online: print(
                f"{ssomib.BROKER_NAME} is now " f"{'online' if online else 'offline'}."
            )
        )
        GLib.MainLoop().run()
        return

    accounts = ssomib.get_accounts()
    if len(accounts["accounts"]) == 0:
        print("warning: no accounts registered.", file=sys.stderr)

    if args.command == "getAccounts":
        json.dump(accounts, indent=2, fp=sys.stdout)
    elif args.command == "getVersion":
        json.dump(ssomib.get_broker_version(), indent=2, fp=sys.stdout)
    elif args.command == "acquirePrtSsoCookie":
        account = _get_account(accounts, args.account)
        cookie = ssomib.acquire_prt_sso_cookie(account, args.ssoUrl)
        json.dump(cookie, indent=2, fp=sys.stdout)
    elif args.command == "acquireTokenSilently":
        account = _get_account(accounts, args.account)
        token = ssomib.acquire_token_silently(account)
        json.dump(token, indent=2, fp=sys.stdout)
    # add newline
    print()


if __name__ == "__main__":
    if "--interactive" in sys.argv or "-i" in sys.argv:
        run_interactive()
    else:
        run_as_native_messaging()
