nest-matrix-bot/matrix_keycloak_bot.py

import asyncio
import os
import sys
import markdown
import json
from keycloak import KeycloakAdmin
from keycloak.exceptions import KeycloakError
from nio import (
    AsyncClient,
    MatrixRoom,
    RoomMessageText,
    ReactionEvent,
    RoomSendResponse,
    InviteMemberEvent,
    OlmUnverifiedDeviceError,
    KeyVerificationStart,
    KeyVerificationKey,
    KeyVerificationMac,
    KeyVerificationCancel,
    ToDeviceError,
    UnknownToDeviceEvent,
    ToDeviceMessage,
)

# ------------------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------------------
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")

CREDENTIALS_FILE = "/app/store/credentials.json"

# In-memory tracking of active requests
pending_requests = {}


# ------------------------------------------------------------------------------
# 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]:
    """Checks if a user is already a member of a specified Keycloak group."""
    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]:
    """Adds a Keycloak user to a target group."""
    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]:
    """Fetches available Keycloak groups, filtering out 'admin' and checking user membership."""
    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)}"


# ------------------------------------------------------------------------------
# Auto-Verification Event Handler
# ------------------------------------------------------------------------------
async def on_to_device(client: AsyncClient, event):
    """Handles incoming interactive SAS key verification events automatically."""
    if isinstance(event, UnknownToDeviceEvent):
        if event.type == "m.key.verification.request":
            content = event.source.get("content", {})
            tx_id = content.get("transaction_id")
            from_device = content.get("from_device", "*")

            if tx_id:
                print(f"Received verification request {tx_id} from {event.sender} ({from_device}). Replying with 'ready'...")
                ready_content = {
                    "from_device": client.device_id,
                    "methods": ["m.sas.v1"],
                    "transaction_id": tx_id,
                }
                
                # Pass the event type, recipient, device, and content payload as separate arguments
                message = ToDeviceMessage(
                    "m.key.verification.ready",
                    event.sender,      # recipient (positional argument 2)
                    from_device,       # recipient_device (positional argument 3)
                    ready_content      # content (positional argument 4)
                )
                
                resp = await client.to_device(message)
                
                if isinstance(resp, ToDeviceError):
                    print(f"Failed to send verification ready: {resp}")
                else:
                    print(f"Sent m.key.verification.ready for {tx_id}. Awaiting start...")

        elif event.type == "m.key.verification.done":
            print(f"✅ Key verification done event received from {event.sender}.")

    elif isinstance(event, KeyVerificationStart):
        print(f"Received key verification start from {event.sender} ({event.from_device}). Accepting start...")
        resp = await client.accept_key_verification(event.transaction_id)
        if isinstance(resp, ToDeviceError):
            print(f"Failed to accept verification start: {resp}")

    elif isinstance(event, KeyVerificationKey):
        print(f"Received SAS key from {event.sender}. Auto-confirming SAS...")
        resp = await client.confirm_short_auth_string(event.transaction_id)
        if isinstance(resp, ToDeviceError):
            print(f"Failed to confirm SAS: {resp}")

    elif isinstance(event, KeyVerificationMac):
        # Extract device_id safely from the raw event payload
        content = event.source.get("content", {})
        from_device = (
            content.get("from_device") 
            or event.source.get("sender_device") 
            or "*"
        )
        
        print(f"Key verification MAC received from {event.sender} ({from_device}). Finalizing verification...")

        # 1. Send m.key.verification.done to complete the SAS exchange
        done_msg = ToDeviceMessage(
            "m.key.verification.done",
            event.sender,
            from_device,
            {"transaction_id": event.transaction_id},
        )
        resp = await client.to_device(done_msg)
        if isinstance(resp, ToDeviceError):
            print(f"Failed to send verification done: {resp}")
        else:
            print(f"Sent m.key.verification.done for {event.transaction_id}.")

        # 2. Mark the sender's device as verified in nio's local device store
        if from_device != "*":
            device = client.device_store.get(event.sender, {}).get(from_device)
            if device:
                client.verify_device(device)
                print(f"✅ Device {from_device} for {event.sender} is now verified and trusted!")
            else:
                print(f"⚠️ Device {from_device} not found in store, but verification completed.")

    elif isinstance(event, KeyVerificationCancel):
        print(f"Key verification cancelled by {event.sender}: {getattr(event, 'reason', 'No reason given')}")


# ------------------------------------------------------------------------------
# Matrix Message Sending Wrappers
# ------------------------------------------------------------------------------
async def send_message_safe(client: AsyncClient, room_id: str, content: dict):
    while True:
        try:
            return await client.room_send(
                room_id=room_id,
                message_type="m.room.message",
                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 send_markdown_message(
    client: AsyncClient,
    room_id: str,
    md_text: str,
    msgtype: str = "m.text",
    relates_to: dict = None,
):
    html_body = markdown.markdown(md_text, extensions=["extra", "sane_lists", "nl2br"])

    content = {
        "msgtype": msgtype,
        "body": md_text,
        "format": "org.matrix.custom.html",
        "formatted_body": html_body,
    }

    if relates_to:
        content["m.relates_to"] = relates_to
        content["m.new_content"] = {
            "msgtype": msgtype,
            "body": md_text,
            "format": "org.matrix.custom.html",
            "formatted_body": html_body,
        }
        content["body"] = f"* {md_text}"
        content["formatted_body"] = f"* {html_body}"

    return await send_message_safe(client, room_id, content)


# ------------------------------------------------------------------------------
# Matrix Event Handlers
# ------------------------------------------------------------------------------
async def on_invite(client: AsyncClient, room: MatrixRoom, event: InviteMemberEvent):
    if event.state_key != client.user_id:
        return
    if not event.sender.endswith(":starling.mynest.love"):
        return

    print(f"Received invite for room {room.room_id} from {event.sender}. Joining...")
    await client.join(room.room_id)


async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageText):
    if event.sender == client.user_id:
        return

    body = event.body.strip()
    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_markdown_message(
            client,
            room_id=room.room_id,
            md_text=msg,
            msgtype="m.notice" if not success else "m.text",
        )
        return

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

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

    requested_group = parts[1].strip()

    if requested_group.lower() == "admin":
        await send_markdown_message(
            client,
            room_id=room.room_id,
            md_text="❌ Requests for the **admin** group are not permitted.",
            msgtype="m.notice",
        )
        return

    # Check if the user is already in the requested group before asking for approval
    already_member, membership_msg = check_group_membership(kc_username, requested_group)
    if already_member or membership_msg:
        await send_markdown_message(
            client,
            room_id=room.room_id,
            md_text=membership_msg,
            msgtype="m.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**."
    )

    res = await send_markdown_message(
        client,
        room_id=ADMIN_ROOM_ID,
        md_text=admin_msg_body,
    )

    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_markdown_message(
            client,
            room_id=room.room_id,
            md_text=f"Request for group `{requested_group}` submitted for admin approval.",
            msgtype="m.notice",
        )


async def on_reaction(client: AsyncClient, room: MatrixRoom, event: ReactionEvent):
    if room.room_id != ADMIN_ROOM_ID or event.sender not in ADMIN_MATRIX_IDS:
        return

    target_event_id = event.reacts_to
    emoji = event.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 `{event.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(
            client,
            room_id=request["origin_room_id"],
            md_text=user_msg,
            msgtype="m.notice",
        )

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

        await send_markdown_message(
            client,
            room_id=request["origin_room_id"],
            md_text=f"Your request for group `{request['requested_group']}` was denied.",
            msgtype="m.notice",
        )
    else:
        return

    # Update original message via m.replace
    await send_markdown_message(
        client,
        room_id=ADMIN_ROOM_ID,
        md_text=status_text,
        relates_to={"rel_type": "m.replace", "event_id": target_event_id},
    )

    del pending_requests[target_event_id]


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

    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
    )
    client.add_to_device_callback(
        lambda event: on_to_device(client, event),
        (
            KeyVerificationStart,
            KeyVerificationKey,
            KeyVerificationMac,
            KeyVerificationCancel,
            UnknownToDeviceEvent,
        ),
    )

    # 1. Attempt to restore an existing session
    if os.path.exists(CREDENTIALS_FILE):
        print("Restoring existing Matrix session...")
        with open(CREDENTIALS_FILE, "r") as f:
            creds = json.load(f)
        
        client.access_token = creds["access_token"]
        client.user_id = creds["user_id"]
        client.device_id = creds["device_id"]
        
    # 2. If no session exists, log in and save the credentials
    else:
        print(f"No existing session found. Logging in to {MATRIX_HOMESERVER}...")
        resp = await client.login(MATRIX_BOT_PASSWORD)
        
        # Create the store directory if it doesn't exist
        os.makedirs(os.path.dirname(CREDENTIALS_FILE), exist_ok=True)
        
        with open(CREDENTIALS_FILE, "w") as f:
            json.dump({
                "access_token": resp.access_token,
                "user_id": resp.user_id,
                "device_id": resp.device_id,
            }, f)
        print(f"Logged in and saved credentials for {MATRIX_BOT_USER}.")

    print("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)