Rewrite with mautrix-python to handle SSSS

This commit is contained in:
Astra Logical 2026-07-23 11:45:09 -04:00
parent d4b9e808e4
commit ec035b5e44
2 changed files with 124 additions and 255 deletions

View file

@ -3,44 +3,44 @@ import os
import sys import sys
import markdown import markdown
import json import json
import logging
from keycloak import KeycloakAdmin from keycloak import KeycloakAdmin
from keycloak.exceptions import KeycloakError from keycloak.exceptions import KeycloakError
from nio import (
AsyncClient, from mautrix.client import Client
MatrixRoom, from mautrix.crypto import OlmMachine
RoomMessageText, from mautrix.crypto.store.sql import SQLCryptoStore
ReactionEvent, from mautrix.client.state_store.sql import SQLStateStore
RoomSendResponse, from mautrix.util.async_db import Database
InviteMemberEvent, from mautrix.types import (
OlmUnverifiedDeviceError, EventType, MessageType, TextMessageEventContent, Format,
KeyVerificationStart, Event, RoomID, Membership, RelationType
KeyVerificationKey,
KeyVerificationMac,
KeyVerificationCancel,
ToDeviceError,
UnknownToDeviceEvent,
ToDeviceMessage,
) )
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
# Configuration # Configuration
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
MATRIX_HOMESERVER = os.getenv("MATRIX_HOMESERVER", "https://matrix.nest.local") MATRIX_HOMESERVER = os.getenv("MATRIX_HOMESERVER", "https://matrix.mynest.love")
MATRIX_BOT_USER = os.getenv("MATRIX_BOT_USER", "@access-bot:nest.local") MATRIX_BOT_USER = os.getenv("MATRIX_BOT_USER", "@access-bot:mynest.love")
MATRIX_BOT_PASSWORD = os.getenv("MATRIX_BOT_PASSWORD", "SuperSecretBotPassword") MATRIX_BOT_PASSWORD = os.getenv("MATRIX_BOT_PASSWORD", "SuperSecretBotPassword")
ADMIN_ROOM_ID = os.getenv("ADMIN_ROOM_ID", "!adminroomid:nest.local") ADMIN_ROOM_ID = os.getenv("ADMIN_ROOM_ID", "!adminroomid:mynest.love")
ADMIN_MATRIX_IDS = set(os.getenv("ADMIN_MATRIX_IDS", "@you:nest.local").split(",")) ADMIN_MATRIX_IDS = set(os.getenv("ADMIN_MATRIX_IDS", "@you:mynest.love").split(","))
MATRIX_RECOVERY_KEY = os.getenv("MATRIX_RECOVERY_KEY", "") # Pass your SSSS string here
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_REALM = os.getenv("KEYCLOAK_REALM", "master")
KEYCLOAK_CLIENT_ID = os.getenv("KEYCLOAK_CLIENT_ID", "nest-matrix-bot") KEYCLOAK_CLIENT_ID = os.getenv("KEYCLOAK_CLIENT_ID", "nest-matrix-bot")
KEYCLOAK_CLIENT_SECRET = os.getenv("KEYCLOAK_CLIENT_SECRET", "your-client-secret") KEYCLOAK_CLIENT_SECRET = os.getenv("KEYCLOAK_CLIENT_SECRET", "your-client-secret")
CREDENTIALS_FILE = "/app/store/credentials.json" CREDENTIALS_FILE = "/app/store/credentials.json"
DB_URI = "sqlite:///app/store/bot_state.db"
# In-memory tracking of active requests # In-memory tracking of active requests (Event ID -> Request Info)
pending_requests = {} pending_requests = {}
logging.basicConfig(level=logging.INFO)
log = logging.getLogger("keycloak_bot")
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
# Keycloak Helper Functions # Keycloak Helper Functions
@ -57,9 +57,7 @@ def get_keycloak_client() -> KeycloakAdmin:
verify=True, verify=True,
) )
def check_group_membership(username: str, group_name: str) -> tuple[bool, str]: 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: try:
kc = get_keycloak_client() kc = get_keycloak_client()
users = kc.get_users(query={"username": username, "exact": True}) users = kc.get_users(query={"username": username, "exact": True})
@ -74,18 +72,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 True, f" You are already a member of the `{group_name}` group."
return False, "" return False, ""
except KeycloakError as e: except KeycloakError as e:
return False, f"⚠️ Keycloak API error: {str(e)}" return False, f"⚠️ Keycloak API error: {str(e)}"
except Exception as e: except Exception as e:
return False, f"⚠️ Unexpected error: {str(e)}" return False, f"⚠️ Unexpected error: {str(e)}"
def add_user_to_kc_group(username: str, group_name: str) -> tuple[bool, str]: def add_user_to_kc_group(username: str, group_name: str) -> tuple[bool, str]:
"""Adds a Keycloak user to a target group."""
try: try:
kc = get_keycloak_client() kc = get_keycloak_client()
users = kc.get_users(query={"username": username, "exact": True}) users = kc.get_users(query={"username": username, "exact": True})
if not users: if not users:
return False, f"Keycloak user '{username}' not found." return False, f"Keycloak user '{username}' not found."
@ -99,18 +93,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) kc.group_user_add(user_id=user_id, group_id=group_id)
return True, f"Successfully added '{username}' to group '{target_group['name']}'." return True, f"Successfully added '{username}' to group '{target_group['name']}'."
except KeycloakError as e: except KeycloakError as e:
return False, f"Keycloak API error: {str(e)}" return False, f"Keycloak API error: {str(e)}"
except Exception as e: except Exception as e:
return False, f"Unexpected error: {str(e)}" return False, f"Unexpected error: {str(e)}"
def get_requestable_groups(username: str) -> tuple[bool, str]: def get_requestable_groups(username: str) -> tuple[bool, str]:
"""Fetches available Keycloak groups, filtering out 'admin' and checking user membership."""
try: try:
kc = get_keycloak_client() kc = get_keycloak_client()
users = kc.get_users(query={"username": username, "exact": True}) users = kc.get_users(query={"username": username, "exact": True})
if not users: if not users:
user_group_names = set() user_group_names = set()
@ -134,7 +124,6 @@ def get_requestable_groups(username: str) -> tuple[bool, str]:
group_list = "\n".join(group_lines) group_list = "\n".join(group_lines)
return True, f"**Available Groups:**\n\n{group_list}" return True, f"**Available Groups:**\n\n{group_list}"
except KeycloakError as e: except KeycloakError as e:
return False, f"Keycloak API error: {str(e)}" return False, f"Keycloak API error: {str(e)}"
except Exception as e: except Exception as e:
@ -142,180 +131,68 @@ def get_requestable_groups(username: str) -> tuple[bool, str]:
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
# Auto-Verification Event Handler # Matrix Helper Functions & Event Handlers
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
async def on_to_device(client: AsyncClient, event): client: Client = None # Global client instance
"""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( async def send_markdown_message(
client: AsyncClient,
room_id: str, room_id: str,
md_text: str, md_text: str,
msgtype: str = "m.text", msgtype: MessageType = MessageType.TEXT,
relates_to: dict = None, 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"]) 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 = { return await client.send_message(room_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)
# ------------------------------------------------------------------------------ async def on_invite(evt: Event):
# Matrix Event Handlers """Auto-joins rooms when invited."""
# ------------------------------------------------------------------------------ if evt.state_key == client.mxid and evt.content.membership == Membership.INVITE:
async def on_invite(client: AsyncClient, room: MatrixRoom, event: InviteMemberEvent): log.info(f"Received invite for room {evt.room_id}. Joining...")
if event.state_key != client.user_id: await client.join_room(evt.room_id)
return
if not event.sender.endswith(":starling.mynest.love"):
async def on_message(evt: Event):
"""Parses incoming text messages for commands."""
if evt.sender == client.mxid:
return return
print(f"Received invite for room {room.room_id} from {event.sender}. Joining...") # Ensure this is a standard text message
await client.join(room.room_id) if not hasattr(evt.content, 'body') or not isinstance(evt.content.body, str):
async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageText):
if event.sender == client.user_id:
return return
body = event.body.strip() body = evt.content.body.strip()
matrix_user = event.sender matrix_user = evt.sender
kc_username = matrix_user.split(":")[0].lstrip("@") kc_username = matrix_user.split(":")[0].lstrip("@")
# --- COMMAND: !groups ---
if body.startswith("!groups"): if body.startswith("!groups"):
success, msg = get_requestable_groups(kc_username) success, msg = get_requestable_groups(kc_username)
await send_markdown_message( await send_markdown_message(
client, room_id=evt.room_id,
room_id=room.room_id,
md_text=msg, md_text=msg,
msgtype="m.notice" if not success else "m.text", msgtype=MessageType.NOTICE if not success else MessageType.TEXT,
) )
return return
# --- COMMAND: !request ---
if not body.startswith("!request"): if not body.startswith("!request"):
return return
parts = body.split(maxsplit=1) parts = body.split(maxsplit=1)
if len(parts) < 2: if len(parts) < 2:
await send_markdown_message( await send_markdown_message(
client, room_id=evt.room_id,
room_id=room.room_id,
md_text="Usage: `!request <group_name>` or `!groups`", md_text="Usage: `!request <group_name>` or `!groups`",
msgtype="m.notice", msgtype=MessageType.NOTICE,
) )
return return
@ -323,21 +200,18 @@ async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageTe
if requested_group.lower() == "admin": if requested_group.lower() == "admin":
await send_markdown_message( await send_markdown_message(
client, room_id=evt.room_id,
room_id=room.room_id,
md_text="❌ Requests for the **admin** group are not permitted.", md_text="❌ Requests for the **admin** group are not permitted.",
msgtype="m.notice", msgtype=MessageType.NOTICE,
) )
return 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) already_member, membership_msg = check_group_membership(kc_username, requested_group)
if already_member or membership_msg: if already_member or membership_msg:
await send_markdown_message( await send_markdown_message(
client, room_id=evt.room_id,
room_id=room.room_id,
md_text=membership_msg, md_text=membership_msg,
msgtype="m.notice", msgtype=MessageType.NOTICE,
) )
return return
@ -348,34 +222,34 @@ async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageTe
f"React with 👍 to **Approve** or 👎 to **Deny**." f"React with 👍 to **Approve** or 👎 to **Deny**."
) )
res = await send_markdown_message( # In mautrix, send_message returns the event ID string directly
client, event_id = await send_markdown_message(
room_id=ADMIN_ROOM_ID, room_id=ADMIN_ROOM_ID,
md_text=admin_msg_body, md_text=admin_msg_body,
) )
if isinstance(res, RoomSendResponse): if event_id:
pending_requests[res.event_id] = { pending_requests[event_id] = {
"matrix_user": matrix_user, "matrix_user": matrix_user,
"kc_username": kc_username, "kc_username": kc_username,
"requested_group": requested_group, "requested_group": requested_group,
"origin_room_id": room.room_id, "origin_room_id": evt.room_id,
} }
await send_markdown_message( await send_markdown_message(
client, room_id=evt.room_id,
room_id=room.room_id,
md_text=f"Request for group `{requested_group}` submitted for admin approval.", 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): async def on_reaction(evt: Event):
if room.room_id != ADMIN_ROOM_ID or event.sender not in ADMIN_MATRIX_IDS: """Handles admin approval/denial via emoji reactions."""
if evt.room_id != ADMIN_ROOM_ID or evt.sender not in ADMIN_MATRIX_IDS:
return return
target_event_id = event.reacts_to target_event_id = evt.content.relates_to.event_id
emoji = event.key emoji = evt.content.relates_to.key
if not target_event_id or target_event_id not in pending_requests: if not target_event_id or target_event_id not in pending_requests:
return return
@ -383,117 +257,111 @@ async def on_reaction(client: AsyncClient, room: MatrixRoom, event: ReactionEven
request = pending_requests[target_event_id] request = pending_requests[target_event_id]
if "👍" in emoji: if "👍" in emoji:
success, msg = add_user_to_kc_group( success, msg = add_user_to_kc_group(request["kc_username"], request["requested_group"])
request["kc_username"], request["requested_group"]
)
status_text = ( 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"**User:** `{request['matrix_user']}` → **Group:** `{request['requested_group']}`\n"
f"**Result:** {msg}" f"**Result:** {msg}"
) )
user_msg = ( user_msg = (
f"🎉 Your request for access to group `{request['requested_group']}` has been approved!" f"🎉 Your request for access to group `{request['requested_group']}` has been approved!"
if success else f"⚠️ Approval failed: {msg}" if success else f"⚠️ Approval failed: {msg}"
) )
await send_markdown_message( await send_markdown_message(
client,
room_id=request["origin_room_id"], room_id=request["origin_room_id"],
md_text=user_msg, md_text=user_msg,
msgtype="m.notice", msgtype=MessageType.NOTICE,
) )
elif "👎" in emoji: elif "👎" in emoji:
status_text = ( 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']}`" f"**User:** `{request['matrix_user']}` → **Group:** `{request['requested_group']}`"
) )
await send_markdown_message( await send_markdown_message(
client,
room_id=request["origin_room_id"], room_id=request["origin_room_id"],
md_text=f"Your request for group `{request['requested_group']}` was denied.", md_text=f"Your request for group `{request['requested_group']}` was denied.",
msgtype="m.notice", msgtype=MessageType.NOTICE,
) )
else: else:
return return
# Update original message via m.replace # Update original message via m.replace
await send_markdown_message( await send_markdown_message(
client,
room_id=ADMIN_ROOM_ID, room_id=ADMIN_ROOM_ID,
md_text=status_text, 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] del pending_requests[target_event_id]
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
# Main Event Loop # Main Execution & Setup
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
async def main(): async def main():
client = AsyncClient( global client
MATRIX_HOMESERVER,
MATRIX_BOT_USER,
store_path="/app/store",
device_id="KEYCLOAK_ACCESS_BOT",
)
client.add_event_callback( # 1. Initialize SQLite backend for stores
lambda room, event: on_message(client, room, event), RoomMessageText db = Database.create(DB_URI)
) await db.start()
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 # 2. Initialize client
client = Client(base_url=MATRIX_HOMESERVER, mxid=MATRIX_BOT_USER)
# 3. Session Persistence
session_restored = False session_restored = False
if os.path.exists(CREDENTIALS_FILE): if os.path.exists(CREDENTIALS_FILE):
print("Restoring existing Matrix session...") log.info("Restoring existing Matrix session...")
try: try:
with open(CREDENTIALS_FILE, "r") as f: with open(CREDENTIALS_FILE, "r") as f:
creds = json.load(f) creds = json.load(f)
client.api.token = creds["access_token"]
client.access_token = creds["access_token"] client.mxid = creds["user_id"]
client.user_id = creds["user_id"]
client.device_id = creds["device_id"] client.device_id = creds["device_id"]
session_restored = True session_restored = True
except (json.JSONDecodeError, KeyError) as e: except (json.JSONDecodeError, KeyError) as e:
print(f"⚠️ Credentials file is empty or corrupted ({e}). Discarding and forcing new login...") log.warning(f"Corrupted credentials file ({e}). Forcing new login...")
# 2. If no session exists or restoration failed, log in and save the credentials
if not session_restored: if not session_restored:
print(f"No valid session found. Logging in to {MATRIX_HOMESERVER}...") log.info(f"Logging in to {MATRIX_HOMESERVER}...")
resp = await client.login(MATRIX_BOT_PASSWORD) resp = await client.login(password=MATRIX_BOT_PASSWORD, device_name="KEYCLOAK_ACCESS_BOT")
# Create the store directory if it doesn't exist
os.makedirs(os.path.dirname(CREDENTIALS_FILE), exist_ok=True) os.makedirs(os.path.dirname(CREDENTIALS_FILE), exist_ok=True)
with open(CREDENTIALS_FILE, "w") as f: with open(CREDENTIALS_FILE, "w") as f:
json.dump({ json.dump({
"access_token": resp.access_token, "access_token": client.api.token,
"user_id": resp.user_id, "user_id": client.mxid,
"device_id": resp.device_id, "device_id": client.device_id,
}, f) }, f)
print(f"Logged in and saved credentials for {MATRIX_BOT_USER}.") log.info("Logged in and saved credentials.")
print("Syncing...") # 4. Attach Crypto & State Stores
await client.sync_forever(timeout=30000, full_state=True) state_store = SQLStateStore(db)
client.state_store = state_store
crypto_store = SQLCryptoStore(db, client.mxid, client.device_id)
machine = OlmMachine(client, crypto_store, state_store)
await machine.load()
client.crypto = machine
# 5. SSSS Recovery Key Ingestion
if MATRIX_RECOVERY_KEY:
try:
log.info("🔑 SSSS Recovery key detected. Attempting to cross-sign session...")
# Instructs mautrix to decrypt your master key from secret storage
await machine.verify_session(MATRIX_RECOVERY_KEY)
log.info("✅ Session successfully self-verified!")
except Exception as e:
log.error(f"⚠️ Could not verify session using recovery key: {e}")
# 6. 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)
# 7. Start Sync Loop
log.info("Starting Mautrix sync loop...")
await client.start()
if __name__ == "__main__": if __name__ == "__main__":

View file

@ -1,3 +1,4 @@
matrix-nio[e2e]>=0.24.0 mautrix[e2e]>=0.20.0
aiosqlite>=0.19.0
python-keycloak>=4.0.0 python-keycloak>=4.0.0
markdown markdown