nest-matrix-bot/matrix_keycloak_bot.py

import asyncio
import os
import sys
import markdown
import json
import logging
from keycloak import KeycloakAdmin
from keycloak.exceptions import KeycloakError

from mautrix.client import Client
from mautrix.crypto import OlmMachine
from mautrix.crypto.store import MemoryCryptoStore
from mautrix.client.state_store import MemoryStateStore
from mautrix.types import (
    EventType, MessageType, TextMessageEventContent, Format,
    Event, RoomID, Membership, RelationType, UserID
)

# ------------------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------------------
MATRIX_HOMESERVER = os.getenv("MATRIX_HOMESERVER", "https://matrix.mynest.love")
MATRIX_BOT_USER = os.getenv("MATRIX_BOT_USER", "@access-bot:mynest.love")
MATRIX_BOT_PASSWORD = os.getenv("MATRIX_BOT_PASSWORD", "SuperSecretBotPassword")
ADMIN_ROOM_ID = os.getenv("ADMIN_ROOM_ID", "!adminroomid:mynest.love")
ADMIN_MATRIX_IDS = set(os.getenv("ADMIN_MATRIX_IDS", "@you:mynest.love").split(","))

KEYCLOAK_URL = os.getenv("KEYCLOAK_URL", "https://keycloak.mynest.love/")
KEYCLOAK_REALM = os.getenv("KEYCLOAK_REALM", "master")
KEYCLOAK_CLIENT_ID = os.getenv("KEYCLOAK_CLIENT_ID", "nest-matrix-bot")
KEYCLOAK_CLIENT_SECRET = os.getenv("KEYCLOAK_CLIENT_SECRET", "your-client-secret")

# In-memory tracking of active requests (Event ID -> Request Info)
pending_requests = {}

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("keycloak_bot")


# ------------------------------------------------------------------------------
# Custom State Store to support find_shared_rooms
# ------------------------------------------------------------------------------
class CustomMemoryStateStore(MemoryStateStore):
    def __init__(self, bot_mxid: str):
        super().__init__()
        self.bot_mxid = bot_mxid

    async def find_shared_rooms(self, user_id: UserID) -> list[RoomID]:
        shared = []
        states = getattr(self, "states", {})
        for room_id in states.keys():
            try:
                rid = RoomID(room_id)
                user_membership = await self.get_membership(rid, user_id)
                bot_membership = await self.get_membership(rid, UserID(self.bot_mxid))
                if user_membership == Membership.JOIN and bot_membership == Membership.JOIN:
                    shared.append(rid)
            except Exception:
                pass
        return shared


# ------------------------------------------------------------------------------
# Keycloak Helper Functions
# ------------------------------------------------------------------------------
def get_keycloak_client() -> KeycloakAdmin:
    """Initializes KeycloakAdmin client using client credentials."""
    return KeycloakAdmin(
        server_url=KEYCLOAK_URL,
        client_id=KEYCLOAK_CLIENT_ID,
        client_secret_key=KEYCLOAK_CLIENT_SECRET,
        realm_name=KEYCLOAK_REALM,
        user_realm_name=KEYCLOAK_REALM,
        grant_type="client_credentials",
        verify=True,
    )

def check_group_membership(username: str, group_name: str) -> tuple[bool, str]:
    try:
        kc = get_keycloak_client()
        users = kc.get_users(query={"username": username, "exact": True})
        if not users:
            return False, f"⚠️ Keycloak user `{username}` not found."

        user_id = users[0]["id"]
        user_groups = kc.get_user_groups(user_id=user_id)
        user_group_names = {g["name"].lower() for g in user_groups}

        if group_name.lower() in user_group_names:
            return True, f"ℹ️ You are already a member of the `{group_name}` group."

        return False, ""
    except KeycloakError as e:
        return False, f"⚠️ Keycloak API error: {str(e)}"
    except Exception as e:
        return False, f"⚠️ Unexpected error: {str(e)}"

def add_user_to_kc_group(username: str, group_name: str) -> tuple[bool, str]:
    try:
        kc = get_keycloak_client()
        users = kc.get_users(query={"username": username, "exact": True})
        if not users:
            return False, f"Keycloak user '{username}' not found."
        user_id = users[0]["id"]

        groups = kc.get_groups()
        target_group = next((g for g in groups if g["name"].lower() == group_name.lower()), None)
        if not target_group:
            return False, f"Keycloak group '{group_name}' not found."
        group_id = target_group["id"]

        kc.group_user_add(user_id=user_id, group_id=group_id)
        return True, f"Successfully added '{username}' to group '{target_group['name']}'."
    except KeycloakError as e:
        return False, f"Keycloak API error: {str(e)}"
    except Exception as e:
        return False, f"Unexpected error: {str(e)}"

def get_requestable_groups(username: str) -> tuple[bool, str]:
    try:
        kc = get_keycloak_client()
        users = kc.get_users(query={"username": username, "exact": True})
        if not users:
            user_group_names = set()
        else:
            user_id = users[0]["id"]
            user_groups = kc.get_user_groups(user_id=user_id)
            user_group_names = {g["name"] for g in user_groups}

        groups = kc.get_groups()
        valid_groups = [g["name"] for g in groups if g["name"].lower() != "admin"]

        if not valid_groups:
            return True, "No groups are currently available to request."

        group_lines = []
        for name in valid_groups:
            if name in user_group_names:
                group_lines.append(f"* `{name}` *(Already a member)*")
            else:
                group_lines.append(f"* `{name}`")

        group_list = "\n".join(group_lines)
        return True, f"**Available Groups:**\n\n{group_list}"
    except KeycloakError as e:
        return False, f"Keycloak API error: {str(e)}"
    except Exception as e:
        return False, f"⚠️ Unexpected error: {str(e)}"


# ------------------------------------------------------------------------------
# Matrix Helper Functions & Event Handlers
# ------------------------------------------------------------------------------
client: Client = None  # Global client instance

async def send_markdown_message(
    room_id: str,
    md_text: str,
    msgtype: MessageType = MessageType.TEXT,
    edit_event_id: str = None,
):
    """Sends a markdown formatted message, optionally replacing an older message."""
    html_body = markdown.markdown(md_text, extensions=["extra", "sane_lists", "nl2br"])
    content = TextMessageEventContent(
        msgtype=msgtype,
        body=md_text,
        format=Format.HTML,
        formatted_body=html_body,
    )
    if edit_event_id:
        content.set_edit(edit_event_id)

    return await client.send_message(room_id, content)


async def on_invite(evt: Event):
    """Auto-joins rooms when invited."""
    if evt.state_key == client.mxid and evt.content.membership == Membership.INVITE:
        log.info(f"Received invite for room {evt.room_id}. Joining...")
        await client.join_room(evt.room_id)


async def on_message(evt: Event):
    """Parses incoming text messages for commands."""
    if evt.sender == client.mxid:
        return

    # Ensure this is a standard text message
    if not hasattr(evt.content, 'body') or not isinstance(evt.content.body, str):
        return

    body = evt.content.body.strip()
    matrix_user = evt.sender
    kc_username = matrix_user.split(":")[0].lstrip("@")

    if body.startswith("!groups"):
        success, msg = get_requestable_groups(kc_username)
        await send_markdown_message(
            room_id=evt.room_id,
            md_text=msg,
            msgtype=MessageType.NOTICE if not success else MessageType.TEXT,
        )
        return

    if not body.startswith("!request"):
        return

    parts = body.split(maxsplit=1)
    if len(parts) < 2:
        await send_markdown_message(
            room_id=evt.room_id,
            md_text="Usage: `!request <group_name>` or `!groups`",
            msgtype=MessageType.NOTICE,
        )
        return

    requested_group = parts[1].strip()

    if requested_group.lower() == "admin":
        await send_markdown_message(
            room_id=evt.room_id,
            md_text="❌ Requests for the **admin** group are not permitted.",
            msgtype=MessageType.NOTICE,
        )
        return

    already_member, membership_msg = check_group_membership(kc_username, requested_group)
    if already_member or membership_msg:
        await send_markdown_message(
            room_id=evt.room_id,
            md_text=membership_msg,
            msgtype=MessageType.NOTICE,
        )
        return

    admin_msg_body = (
        f"📋 **Access Request**\n\n"
        f"**User:** `{matrix_user}` (KC: `{kc_username}`)\n"
        f"**Requested Group:** `{requested_group}`\n\n"
        f"React with 👍 to **Approve** or 👎 to **Deny**."
    )

    event_id = await send_markdown_message(
        room_id=ADMIN_ROOM_ID,
        md_text=admin_msg_body,
    )

    if event_id:
        pending_requests[event_id] = {
            "matrix_user": matrix_user,
            "kc_username": kc_username,
            "requested_group": requested_group,
            "origin_room_id": evt.room_id,
        }

        await send_markdown_message(
            room_id=evt.room_id,
            md_text=f"Request for group `{requested_group}` submitted for admin approval.",
            msgtype=MessageType.NOTICE,
        )


async def on_reaction(evt: Event):
    """Handles admin approval/denial via emoji reactions."""
    if evt.room_id != ADMIN_ROOM_ID or evt.sender not in ADMIN_MATRIX_IDS:
        return

    target_event_id = evt.content.relates_to.event_id
    emoji = evt.content.relates_to.key

    if not target_event_id or target_event_id not in pending_requests:
        return

    request = pending_requests[target_event_id]

    if "👍" in emoji:
        success, msg = add_user_to_kc_group(request["kc_username"], request["requested_group"])
        status_text = (
            f"✅ **APPROVED** by `{evt.sender}`\n\n"
            f"**User:** `{request['matrix_user']}` → **Group:** `{request['requested_group']}`\n"
            f"**Result:** {msg}"
        )
        user_msg = (
            f"🎉 Your request for access to group `{request['requested_group']}` has been approved!"
            if success else f"⚠️ Approval failed: {msg}"
        )
        await send_markdown_message(
            room_id=request["origin_room_id"],
            md_text=user_msg,
            msgtype=MessageType.NOTICE,
        )

    elif "👎" in emoji:
        status_text = (
            f"❌ **DENIED** by `{evt.sender}`\n\n"
            f"**User:** `{request['matrix_user']}` → **Group:** `{request['requested_group']}`"
        )
        await send_markdown_message(
            room_id=request["origin_room_id"],
            md_text=f"Your request for group `{request['requested_group']}` was denied.",
            msgtype=MessageType.NOTICE,
        )
    else:
        return

    await send_markdown_message(
        room_id=ADMIN_ROOM_ID,
        md_text=status_text,
        edit_event_id=target_event_id,
    )
    del pending_requests[target_event_id]


# ------------------------------------------------------------------------------
# Main Execution & Setup
# ------------------------------------------------------------------------------
async def main():
    global client

    # 1. Initialize client
    client = Client(base_url=MATRIX_HOMESERVER, mxid=MATRIX_BOT_USER)

    # 2. Attach Custom Memory Crypto & State Stores first so keys align with login
    state_store = CustomMemoryStateStore(client.mxid)
    client.state_store = state_store

    crypto_store = MemoryCryptoStore(client.mxid, client.device_id)
    machine = OlmMachine(client, crypto_store, state_store)
    await machine.load()
    client.crypto = machine

    # 3. Perform Fresh Login on Startup (Syncs device ID and crypto keys correctly)
    log.info(f"Logging in to {MATRIX_HOMESERVER}...")
    await client.login(password=MATRIX_BOT_PASSWORD, device_name="KEYCLOAK_ACCESS_BOT")
    log.info("Successfully logged in with fresh device session.")

    # 4. Register Event Handlers
    client.add_event_handler(EventType.ROOM_MEMBER, on_invite)
    client.add_event_handler(EventType.ROOM_MESSAGE, on_message)
    client.add_event_handler(EventType.REACTION, on_reaction)

    # 5. Start Sync Loop
    log.info("Starting Mautrix sync loop...")
    await client.start(None)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nBot stopped by user.")
        sys.exit(0)