Impliment find_shared_rooms

This commit is contained in:
Astra Logical 2026-07-23 12:01:14 -04:00
parent b1fee3e4ca
commit 8cf971c589

View file

@ -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)