From 8cf971c589b06de7ed161e423b66909e712b99c5 Mon Sep 17 00:00:00 2001
From: Astra Logical <booksandpi@gmail.com>
Date: Thu, 23 Jul 2026 12:01:14 -0400
Subject: [PATCH] Impliment find_shared_rooms

---
 matrix_keycloak_bot.py | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/matrix_keycloak_bot.py b/matrix_keycloak_bot.py
index 29e7764..ac66234 100644
--- a/matrix_keycloak_bot.py
+++ b/matrix_keycloak_bot.py
@@ -13,7 +13,7 @@ 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
+    Event, RoomID, Membership, RelationType, UserID
 )
 
 # ------------------------------------------------------------------------------
@@ -39,6 +39,29 @@ 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
 # ------------------------------------------------------------------------------
@@ -326,8 +349,8 @@ async def main():
             }, f)
         log.info("Logged in and saved credentials.")
 
-    # 3. Attach Memory Crypto & State Stores
-    state_store = MemoryStateStore()
+    # 3. Attach Custom Memory Crypto & State Stores
+    state_store = CustomMemoryStateStore(client.mxid)
     client.state_store = state_store
 
     crypto_store = MemoryCryptoStore(client.mxid, client.device_id)