nest-matrix-bot/matrix_keycloak_bot.py

import asyncio
import os
import sys
from keycloak import KeycloakAdmin
from keycloak.exceptions import KeycloakError
from nio import (
    AsyncClient,
    MatrixRoom,
    RoomMessageText,
    ReactionEvent,
    RoomSendResponse,
    InviteMemberEvent,
    OlmUnverifiedDeviceError
)

# ------------------------------------------------------------------------------
# Configuration (Set these via environment variables or replace inline)
# ------------------------------------------------------------------------------
MATRIX_HOMESERVER = os.getenv("MATRIX_HOMESERVER", "https://matrix.nest.local")
MATRIX_BOT_USER = os.getenv("MATRIX_BOT_USER", "@access-bot:nest.local")
MATRIX_BOT_PASSWORD = os.getenv("MATRIX_BOT_PASSWORD", "SuperSecretBotPassword")
ADMIN_ROOM_ID = os.getenv("ADMIN_ROOM_ID", "!adminroomid:nest.local")
ADMIN_MATRIX_IDS = set(os.getenv("ADMIN_MATRIX_IDS", "@you:nest.local").split(","))

KEYCLOAK_URL = os.getenv("KEYCLOAK_URL", "https://keycloak.nest.local/")
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 State
# Maps target message event_id -> request metadata
# ------------------------------------------------------------------------------
pending_requests = {}


# ------------------------------------------------------------------------------
# Keycloak Helper Functions
# ------------------------------------------------------------------------------
def get_keycloak_client() -> KeycloakAdmin:
    """Initializes and returns a KeycloakAdmin client using service account 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 get_requestable_groups(username: str) -> tuple[bool, str]:
    """Fetches available Keycloak groups, filtering out the 'admin' group, and checks user membership."""
    try:
        kc = get_keycloak_client()
        
        # 1. Fetch User UUID to get their current groups
        users = kc.get_users(query={"username": username, "exact": True})
        if not users:
            # If the user doesn't exist in Keycloak, we just show groups normally
            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}

        # 2. Fetch all groups
        groups = kc.get_groups()
        
        # Filter out 'admin' (case-insensitive)
        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."
            
        # 3. Format the output list with membership status
        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{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)}"

def add_user_to_kc_group(username: str, group_name: str) -> tuple[bool, str]:
    """Queries Keycloak for the user and group UUIDs and adds the user to the group."""
    try:
        kc = get_keycloak_client()

        # 1. Fetch User UUID (Assumes Matrix localpart matches Keycloak username)
        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"]

        # 2. Fetch Group UUID
        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"]

        # 3. Assign User to Group
        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)}"


# ------------------------------------------------------------------------------
# Matrix Event Handlers
# ------------------------------------------------------------------------------
async def on_invite(client: AsyncClient, room: MatrixRoom, event: InviteMemberEvent):
    """Auto-joins rooms when invited by an authorized user/homeserver."""
    # Ensure the invite is actually meant for this bot
    if event.state_key != client.user_id:
        return

    # Only accept invites from my domain
    if not event.sender.endswith(":starling.mynest.love"): return

    print(f"Received invite for room {room.room_id} from {event.sender}. Joining...")
    
    response = await client.join(room.room_id)
    if hasattr(response, "room_id"):
        print(f"Successfully joined room {room.room_id}")
    else:
        print(f"Failed to join room {room.room_id}: {response}")


async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageText):
    """Listens for commands."""
    # Ignore messages sent by the bot itself
    if event.sender == client.user_id:
        return

    body = event.body.strip()
    
    # Extract username early so both commands can utilize it
    matrix_user = event.sender
    kc_username = matrix_user.split(":")[0].lstrip("@")
    
    # --- COMMAND: !groups ---
    if body.startswith("!groups"):
        success, msg = get_requestable_groups(kc_username)
        await send_message_safe(
            client,
            room_id=room.room_id,
            content={
                "msgtype": "m.notice" if not success else "m.text",
                "body": msg,
                "format": "org.matrix.custom.html",
                "formatted_body": msg.replace("\n", "<br>")
            },
        )
        return

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

    parts = body.split(maxsplit=1)
    if len(parts) < 2:
        await send_message_safe(
            client,
            room_id=room.room_id,
            content={"msgtype": "m.notice", "body": "Usage: !request <group_name> or !groups"},
        )
        return

    requested_group = parts[1].strip()
    
    # Block 'admin' requests immediately
    if requested_group.lower() == "admin":
        await send_message_safe(
            client,
            room_id=room.room_id,
            content={
                "msgtype": "m.notice", 
                "body": "❌ Requests for the 'admin' group are not permitted."
            },
        )
        return

    matrix_user = event.sender
    kc_username = matrix_user.split(":")[0].lstrip("@")

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

    res = await send_message_safe(
        client,
        room_id=ADMIN_ROOM_ID,
        content={
            "msgtype": "m.text",
            "body": admin_msg_body,
            "format": "org.matrix.custom.html",
            "formatted_body": admin_msg_body.replace("\n", "<br>"),
        },
    )

    if isinstance(res, RoomSendResponse):
        pending_requests[res.event_id] = {
            "matrix_user": matrix_user,
            "kc_username": kc_username,
            "requested_group": requested_group,
            "origin_room_id": room.room_id,
        }

        await send_message_safe(
            client,
            room_id=room.room_id,
            content={
                "msgtype": "m.notice",
                "body": f"Request for group '{requested_group}' submitted for admin approval.",
            },
        )

async def send_message_safe(
    client: AsyncClient, 
    room_id: str, 
    content: dict, 
    message_type: str = "m.room.message"
):
    """Sends a message to a room, automatically trusting any unverified devices if an OlmUnverifiedDeviceError occurs."""
    while True:
        try:
            return await client.room_send(
                room_id=room_id,
                message_type=message_type,
                content=content,
            )
        except OlmUnverifiedDeviceError as e:
            print(f"Auto-trusting unverified device '{e.device.device_id}' for user '{e.device.user_id}'...")
            client.verify_device(e.device)

async def on_reaction(client: AsyncClient, room: MatrixRoom, event: ReactionEvent):
    """Listens for reaction approvals/denials in the admin room."""
    print(f"[REACTION] Received '{event.key}' from {event.sender} on event '{event.reacts_to}'")

    if room.room_id != ADMIN_ROOM_ID:
        return

    # Verify reactor is an authorized admin
    if event.sender not in ADMIN_MATRIX_IDS:
        print(f"[REACTION IGNORED] Sender '{event.sender}' is not in ADMIN_MATRIX_IDS.")
        return

    target_event_id = event.reacts_to
    emoji = event.key

    # Check if this reaction targets an active pending request
    if not target_event_id or target_event_id not in pending_requests:
        print(f"[REACTION IGNORED] Event ID '{target_event_id}' not found in pending_requests.")
        return

    request = pending_requests[target_event_id]

    # Use substring check to catch Unicode variation selectors (e.g. 👍 vs 👍\ufe0f)
    if "👍" in emoji:
        # Execute Keycloak Group Addition
        success, msg = add_user_to_kc_group(
            request["kc_username"], request["requested_group"]
        )

        status_text = (
            f"✅ **APPROVED** by `{event.sender}`\n"
            f"**User:** `{request['matrix_user']}` → **Group:** `{request['requested_group']}`\n"
            f"**Result:** {msg}"
        )

        # Notify requesting user
        user_msg = (
            f"🎉 Your request for access to group '{request['requested_group']}' "
            f"has been approved!" if success else f"⚠️ Approval failed: {msg}"
        )
        await send_message_safe(
            client,
            room_id=request["origin_room_id"],
            content={"msgtype": "m.notice", "body": user_msg},
        )

    elif "👎" in emoji:
        status_text = (
            f"❌ **DENIED** by `{event.sender}`\n"
            f"**User:** `{request['matrix_user']}` → **Group:** `{request['requested_group']}`"
        )

        await send_message_safe(
            client,
            room_id=request["origin_room_id"],
            content={
                "msgtype": "m.notice",
                "body": f"Your request for group '{request['requested_group']}' was denied.",
            },
        )
    else:
        print(f"[REACTION IGNORED] Emoji '{emoji}' is neither 👍 nor 👎.")
        return  # Ignore other reactions

    # Update the original admin room message to reflect resolution status
    fallback_text = f"* {status_text}"
    
    await send_message_safe(
        client,
        room_id=ADMIN_ROOM_ID,
        content={
            "msgtype": "m.text",
            "body": fallback_text,
            "format": "org.matrix.custom.html",
            "formatted_body": fallback_text.replace("\n", "<br>"),
            "m.new_content": {
                "msgtype": "m.text",
                "body": status_text,
                "format": "org.matrix.custom.html",
                "formatted_body": status_text.replace("\n", "<br>"),
            },
            "m.relates_to": {
                "rel_type": "m.replace", 
                "event_id": target_event_id
            },
        },
    )

    # Clean up request state
    del pending_requests[target_event_id]


# ------------------------------------------------------------------------------
# Main Event Loop
# ------------------------------------------------------------------------------
async def main():
    client = AsyncClient(
        MATRIX_HOMESERVER, 
        MATRIX_BOT_USER, 
        store_path="/app/store"
    )

    # Register callbacks
    client.add_event_callback(
        lambda room, event: on_message(client, room, event), RoomMessageText
    )
    client.add_event_callback(
        lambda room, event: on_reaction(client, room, event), ReactionEvent
    )
    client.add_event_callback(
        lambda room, event: on_invite(client, room, event), InviteMemberEvent
    )

    print(f"Connecting to Matrix homeserver {MATRIX_HOMESERVER}...")
    await client.login(MATRIX_BOT_PASSWORD)
    print(f"Logged in as {MATRIX_BOT_USER}. Syncing...")

    await client.sync_forever(timeout=30000, full_state=True)


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