Transition to Mautrix to enable SSSS

This commit is contained in:
Astra Logical 2026-07-23 15:57:26 -04:00
parent d4b9e808e4
commit 0e9ff14c27
3 changed files with 171 additions and 273 deletions

View file

@ -5,10 +5,12 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
WORKDIR /app
# Install libolm C library and build dependencies for E2E crypto bindings
# Install libolm C library, CMake, C++ compiler, and Python headers 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,44 +3,70 @@ import os
import sys
import markdown
import json
import pickle
import logging
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,
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 mautrix.util.async_db import Database
from mautrix.crypto.store.asyncpg import PgCryptoStore
# ------------------------------------------------------------------------------
# Configuration
# ------------------------------------------------------------------------------
MATRIX_HOMESERVER = os.getenv("MATRIX_HOMESERVER", "https://matrix.nest.local")
MATRIX_BOT_USER = os.getenv("MATRIX_BOT_USER", "@access-bot:nest.local")
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:nest.local")
ADMIN_MATRIX_IDS = set(os.getenv("ADMIN_MATRIX_IDS", "@you:nest.local").split(","))
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.nest.local/")
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")
CREDENTIALS_FILE = "/app/store/credentials.json"
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")
# In-memory tracking of active requests
# In-memory tracking of active requests (Event ID -> Request Info)
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
@ -57,9 +83,7 @@ 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})
@ -74,18 +98,14 @@ 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."
@ -99,18 +119,14 @@ 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()
@ -134,188 +150,74 @@ 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)}"
# ------------------------------------------------------------------------------
# Auto-Verification Event Handler
# Matrix Helper Functions & Event Handlers
# ------------------------------------------------------------------------------
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)
client: Client = None # Global client instance
async def send_markdown_message(
client: AsyncClient,
room_id: str,
md_text: str,
msgtype: str = "m.text",
relates_to: dict = None,
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)
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)
return await client.send_message(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"):
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
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:
if not hasattr(evt.content, 'body') or not isinstance(evt.content.body, str):
return
body = event.body.strip()
matrix_user = event.sender
body = evt.content.body.strip()
matrix_user = evt.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,
room_id=evt.room_id,
md_text=msg,
msgtype="m.notice" if not success else "m.text",
msgtype=MessageType.NOTICE if not success else MessageType.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,
room_id=evt.room_id,
md_text="Usage: `!request <group_name>` or `!groups`",
msgtype="m.notice",
msgtype=MessageType.NOTICE,
)
return
@ -323,21 +225,18 @@ async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageTe
if requested_group.lower() == "admin":
await send_markdown_message(
client,
room_id=room.room_id,
room_id=evt.room_id,
md_text="❌ Requests for the **admin** group are not permitted.",
msgtype="m.notice",
msgtype=MessageType.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,
room_id=evt.room_id,
md_text=membership_msg,
msgtype="m.notice",
msgtype=MessageType.NOTICE,
)
return
@ -348,34 +247,33 @@ async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageTe
f"React with 👍 to **Approve** or 👎 to **Deny**."
)
res = await send_markdown_message(
client,
event_id = await send_markdown_message(
room_id=ADMIN_ROOM_ID,
md_text=admin_msg_body,
)
if isinstance(res, RoomSendResponse):
pending_requests[res.event_id] = {
if event_id:
pending_requests[event_id] = {
"matrix_user": matrix_user,
"kc_username": kc_username,
"requested_group": requested_group,
"origin_room_id": room.room_id,
"origin_room_id": evt.room_id,
}
await send_markdown_message(
client,
room_id=room.room_id,
room_id=evt.room_id,
md_text=f"Request for group `{requested_group}` submitted for admin approval.",
msgtype="m.notice",
msgtype=MessageType.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:
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 = event.reacts_to
emoji = event.key
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
@ -383,117 +281,112 @@ async def on_reaction(client: AsyncClient, room: MatrixRoom, event: ReactionEven
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 `{event.sender}`\n\n"
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(
client,
room_id=request["origin_room_id"],
md_text=user_msg,
msgtype="m.notice",
msgtype=MessageType.NOTICE,
)
elif "👎" in emoji:
status_text = (
f"❌ **DENIED** by `{event.sender}`\n\n"
f"❌ **DENIED** by `{evt.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",
msgtype=MessageType.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},
edit_event_id=target_event_id,
)
del pending_requests[target_event_id]
# ------------------------------------------------------------------------------
# Main Event Loop
# Main Execution & Setup
# ------------------------------------------------------------------------------
async def main():
client = AsyncClient(
MATRIX_HOMESERVER,
MATRIX_BOT_USER,
store_path="/app/store",
device_id="KEYCLOAK_ACCESS_BOT",
)
global client
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...")
# Load session if it exists
session_data = {}
if os.path.exists(SESSION_FILE):
try:
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"]
session_restored = True
except (json.JSONDecodeError, KeyError) as e:
print(f"⚠️ Credentials file is empty or corrupted ({e}). Discarding and forcing new login...")
# 2. If no session exists or restoration failed, log in and save the credentials
if not session_restored:
print(f"No valid 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}.")
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
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")
)
# 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. 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
print("Syncing...")
await client.sync_forever(timeout=30000, full_state=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:
json.dump({
"access_token": client.api.token,
"device_id": client.device_id
}, f)
# 3. Initialize the OlmMachine and explicitly LOAD it
client.crypto = OlmMachine(client, crypto_store, state_store)
await client.crypto.load()
# 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__":

View file

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