Compare commits

..

2 commits

Author SHA1 Message Date
Astra Logical
6afa75b638 Fix syncing issues 2026-07-30 12:14:57 -04:00
Astra Logical
195aeff39e Explicitly call load store 2026-07-30 12:08:43 -04:00
3 changed files with 275 additions and 168 deletions

View file

@ -5,12 +5,10 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
WORKDIR /app
# Install libolm C library, CMake, C++ compiler, and Python headers for E2E crypto bindings
# Install libolm C library and build dependencies for E2E crypto bindings
RUN apt-get update && apt-get install -y --no-install-recommends \
libolm-dev \
gcc \
g++ \
cmake \
python3-dev \
&& rm -rf /var/lib/apt/lists/*

View file

@ -3,70 +3,44 @@ 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
from nio import (
AsyncClient,
MatrixRoom,
RoomMessageText,
ReactionEvent,
RoomSendResponse,
InviteMemberEvent,
OlmUnverifiedDeviceError,
KeyVerificationStart,
KeyVerificationKey,
KeyVerificationMac,
KeyVerificationCancel,
ToDeviceError,
UnknownToDeviceEvent,
ToDeviceMessage,
)
from mautrix.util.async_db import Database
from mautrix.crypto.store.asyncpg import PgCryptoStore
# ------------------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------------------
MATRIX_HOMESERVER = os.getenv("MATRIX_HOMESERVER", "https://matrix.mynest.love")
MATRIX_BOT_USER = os.getenv("MATRIX_BOT_USER", "@access-bot:mynest.love")
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:mynest.love")
ADMIN_MATRIX_IDS = set(os.getenv("ADMIN_MATRIX_IDS", "@you:mynest.love").split(","))
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.mynest.love/")
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")
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")
CRYPTO_PICKLE_KEY = os.getenv("CRYPTO_PICKLE_KEY", "secret-passphrase-for-olm-store")
CREDENTIALS_FILE = "/app/store/credentials.json"
# In-memory tracking of active requests (Event ID -> Request Info)
# In-memory tracking of active requests
pending_requests = {}
logging.basicConfig(level=logging.INFO)
log = logging.getLogger("keycloak_bot")
# ------------------------------------------------------------------------------
# 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
@ -83,7 +57,9 @@ def get_keycloak_client() -> KeycloakAdmin:
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})
@ -98,14 +74,18 @@ def check_group_membership(username: str, group_name: str) -> tuple[bool, str]:
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."
@ -119,14 +99,18 @@ def add_user_to_kc_group(username: str, group_name: str) -> tuple[bool, str]:
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()
@ -150,74 +134,188 @@ def get_requestable_groups(username: str) -> tuple[bool, str]:
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)}"
return False, f"Unexpected error: {str(e)}"
# ------------------------------------------------------------------------------
# Matrix Helper Functions & Event Handlers
# Auto-Verification Event Handler
# ------------------------------------------------------------------------------
client: Client = None # Global client instance
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: MessageType = MessageType.TEXT,
edit_event_id: str = None,
msgtype: str = "m.text",
relates_to: dict = 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)
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)
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:
# ------------------------------------------------------------------------------
# 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
if not hasattr(evt.content, 'body') or not isinstance(evt.content.body, str):
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 = evt.content.body.strip()
matrix_user = evt.sender
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(
room_id=evt.room_id,
client,
room_id=room.room_id,
md_text=msg,
msgtype=MessageType.NOTICE if not success else MessageType.TEXT,
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(
room_id=evt.room_id,
client,
room_id=room.room_id,
md_text="Usage: `!request <group_name>` or `!groups`",
msgtype=MessageType.NOTICE,
msgtype="m.notice",
)
return
@ -225,18 +323,21 @@ async def on_message(evt: Event):
if requested_group.lower() == "admin":
await send_markdown_message(
room_id=evt.room_id,
client,
room_id=room.room_id,
md_text="❌ Requests for the **admin** group are not permitted.",
msgtype=MessageType.NOTICE,
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(
room_id=evt.room_id,
client,
room_id=room.room_id,
md_text=membership_msg,
msgtype=MessageType.NOTICE,
msgtype="m.notice",
)
return
@ -247,33 +348,34 @@ async def on_message(evt: Event):
f"React with 👍 to **Approve** or 👎 to **Deny**."
)
event_id = await send_markdown_message(
res = await send_markdown_message(
client,
room_id=ADMIN_ROOM_ID,
md_text=admin_msg_body,
)
if event_id:
pending_requests[event_id] = {
if isinstance(res, RoomSendResponse):
pending_requests[res.event_id] = {
"matrix_user": matrix_user,
"kc_username": kc_username,
"requested_group": requested_group,
"origin_room_id": evt.room_id,
"origin_room_id": room.room_id,
}
await send_markdown_message(
room_id=evt.room_id,
client,
room_id=room.room_id,
md_text=f"Request for group `{requested_group}` submitted for admin approval.",
msgtype=MessageType.NOTICE,
msgtype="m.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:
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 = evt.content.relates_to.event_id
emoji = evt.content.relates_to.key
target_event_id = event.reacts_to
emoji = event.key
if not target_event_id or target_event_id not in pending_requests:
return
@ -281,112 +383,122 @@ async def on_reaction(evt: Event):
request = pending_requests[target_event_id]
if "👍" in emoji:
success, msg = add_user_to_kc_group(request["kc_username"], request["requested_group"])
success, msg = add_user_to_kc_group(
request["kc_username"], request["requested_group"]
)
status_text = (
f"✅ **APPROVED** by `{evt.sender}`\n\n"
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=MessageType.NOTICE,
msgtype="m.notice",
)
elif "👎" in emoji:
status_text = (
f"❌ **DENIED** by `{evt.sender}`\n\n"
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=MessageType.NOTICE,
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,
edit_event_id=target_event_id,
relates_to={"rel_type": "m.replace", "event_id": target_event_id},
)
del pending_requests[target_event_id]
# ------------------------------------------------------------------------------
# Main Execution & Setup
# Main Event Loop
# ------------------------------------------------------------------------------
async def main():
global client
client = AsyncClient(
MATRIX_HOMESERVER,
MATRIX_BOT_USER,
store_path="/app/store",
device_id="KEYCLOAK_ACCESS_BOT",
)
# Load session if it exists
session_data = {}
if os.path.exists(SESSION_FILE):
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
session_restored = False
if os.path.exists(CREDENTIALS_FILE):
print("Restoring existing Matrix session...")
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}")
with open(CREDENTIALS_FILE, "r") as f:
creds = json.load(f)
# Initialize client
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")
)
client.access_token = creds["access_token"]
client.user_id = creds["user_id"]
client.device_id = creds["device_id"]
client.load_store()
session_restored = True
except (json.JSONDecodeError, KeyError) as e:
print(f"⚠️ Credentials file is empty or corrupted ({e}). Discarding and forcing new login...")
# 1. Initialize Mautrix's built-in Database wrapper
db_path = CRYPTO_STORE_FILE.replace(".pkl", ".db")
db = Database.create(
f"sqlite:///{db_path}",
upgrade_table=PgCryptoStore.upgrade_table
)
await db.start()
# 2. If no session exists or restoration failed, log in and save credentials
if not session_restored:
print(f"No valid session found. Logging in to {MATRIX_HOMESERVER}...")
resp = await client.login(MATRIX_BOT_PASSWORD)
# 2. Instantiate PgCryptoStore with account_id, pickle_key, and db
crypto_store = PgCryptoStore(
account_id=client.mxid,
pickle_key=CRYPTO_PICKLE_KEY,
db=db,
)
state_store = CustomMemoryStateStore(client.mxid)
client.state_store = state_store
os.makedirs(os.path.dirname(CREDENTIALS_FILE), exist_ok=True)
# 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"
)
session_dir = os.path.dirname(SESSION_FILE)
if session_dir:
os.makedirs(session_dir, exist_ok=True)
with open(SESSION_FILE, "w") as f:
with open(CREDENTIALS_FILE, "w") as f:
json.dump({
"access_token": client.api.token,
"device_id": client.device_id
"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}.")
# 3. Initialize the OlmMachine and explicitly LOAD it
client.crypto = OlmMachine(client, crypto_store, state_store)
await client.crypto.load()
# Optional: Perform an initial full state sync ONCE on startup
print("Performing initial state sync...")
await client.sync(full_state=True)
# 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)
# Start the continuous sync loop WITHOUT full_state=True and with a 20s timeout
print("Syncing...")
await client.sync_forever(timeout=20000, full_state=False)
if __name__ == "__main__":

View file

@ -1,6 +1,3 @@
mautrix[encryption]>=0.20.0
python-olm>=3.2.0
aiosqlite>=0.19.0
asyncpg>=0.29.0
matrix-nio[e2e]>=0.24.0
python-keycloak>=4.0.0
markdown