import asyncio import os import sys import markdown import json import pickle import logging from keycloak import KeycloakAdmin from keycloak.exceptions import KeycloakError from mautrix.client import Client from mautrix.crypto import OlmMachine from mautrix.crypto.account import OlmAccount 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") SESSION_FILE = os.getenv("MATRIX_SESSION_FILE", "/app/store/bot_session.json") CRYPTO_STORE_FILE = os.getenv("MATRIX_CRYPTO_STORE", "/app/store/crypto_store.pkl") # In-memory tracking of active requests (Event ID -> Request Info) pending_requests = {} logging.basicConfig(level=logging.INFO) log = logging.getLogger("keycloak_bot") # ------------------------------------------------------------------------------ # Custom Pickle-Backed Crypto Store for Persistence # ------------------------------------------------------------------------------ class PickleCryptoStore(MemoryCryptoStore): def __init__(self, filepath: str, account_id: str = "bot", pickle_key: str = ""): super().__init__(account_id=account_id, pickle_key=pickle_key) self.filepath = filepath self.load() def load(self): if os.path.exists(self.filepath): try: with open(self.filepath, "rb") as f: data = pickle.load(f) self.accounts = data.get("accounts", {}) self.sessions = data.get("sessions", {}) self.inbound_group_sessions = data.get("inbound_group_sessions", {}) self.outbound_group_sessions = data.get("outbound_group_sessions", {}) self.devices = data.get("devices", {}) self.cross_signing_keys = data.get("cross_signing_keys", {}) log.info(f"Loaded crypto store from {self.filepath}") except Exception as e: log.warning(f"Failed to load pickle crypto store: {e}") def save(self): try: data = { "accounts": self.accounts, "sessions": self.sessions, "inbound_group_sessions": self.inbound_group_sessions, "outbound_group_sessions": self.outbound_group_sessions, "devices": self.devices, "cross_signing_keys": self.cross_signing_keys, } directory = os.path.dirname(self.filepath) if directory: os.makedirs(directory, exist_ok=True) with open(self.filepath, "wb") as f: pickle.dump(data, f) except Exception as e: log.error(f"Failed to save pickle crypto store: {e}") async def put_account(self, *args, **kwargs): await super().put_account(*args, **kwargs) self.save() async def put_session(self, *args, **kwargs): await super().put_session(*args, **kwargs) self.save() async def put_inbound_group_session(self, *args, **kwargs): await super().put_inbound_group_session(*args, **kwargs) self.save() async def put_outbound_group_session(self, *args, **kwargs): await super().put_outbound_group_session(*args, **kwargs) self.save() async def put_device(self, *args, **kwargs): await super().put_device(*args, **kwargs) self.save() async def put_cross_signing_key(self, user_id, usage, key): await super().put_cross_signing_key(user_id, usage, key) self.save() # ------------------------------------------------------------------------------ # Custom State Store supporting 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 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 ` 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 # Load session if it exists session_data = {} if os.path.exists(SESSION_FILE): try: with open(SESSION_FILE, "r") as f: session_data = json.load(f) except Exception as e: log.warning(f"Could not load session file: {e}") # Initialize client with persisted token and device ID if available client = Client( base_url=MATRIX_HOMESERVER, mxid=MATRIX_BOT_USER, token=session_data.get("access_token"), device_id=session_data.get("device_id", "KEYCLOAK_BOT_DEVICE") ) # Initialize stores using Pickle and Custom Memory State Stores state_store = CustomMemoryStateStore(client.mxid) crypto_store = PickleCryptoStore(CRYPTO_STORE_FILE, account_id=client.mxid) client.state_store = state_store client.crypto = OlmMachine(client, crypto_store, state_store) # Authenticate if access token is missing if not client.api.token: log.info("Logging in to Matrix homeserver...") await client.login( password=MATRIX_BOT_PASSWORD, device_name="KEYCLOAK_ACCESS_BOT", device_id="KEYCLOAK_BOT_DEVICE" ) # Save session data for future restarts session_dir = os.path.dirname(SESSION_FILE) if session_dir: os.makedirs(session_dir, exist_ok=True) with open(SESSION_FILE, "w") as f: json.dump({ "access_token": client.api.token, "device_id": client.device_id }, f) # Ensure an Olm crypto account exists in the store account = await crypto_store.get_account() if not account: log.info("Creating new Olm account...") account = OlmAccount() await crypto_store.put_account(account) # 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) # 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)