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 markdown
import json
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.store.sql import SQLCryptoStore
from mautrix.client.state_store.sql import SQLStateStore
from mautrix.util.async_db import Database
from mautrix.types import (
EventType, MessageType, TextMessageEventContent, Format,
Event, RoomID, Membership, RelationType
)
# ------------------------------------------------------------------------------
# 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(","))
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_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"
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 = {}
logging.basicConfig(level=logging.INFO)
log = logging.getLogger("keycloak_bot")
# ------------------------------------------------------------------------------
# Keycloak Helper Functions
@ -57,9 +57,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 +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 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 +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)
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,7 +124,6 @@ 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:
@ -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):
"""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:
# Ensure this is a standard text message
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 +200,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 +222,34 @@ async def on_message(client: AsyncClient, room: MatrixRoom, event: RoomMessageTe
f"React with 👍 to **Approve** or 👎 to **Deny**."
)
res = await send_markdown_message(
client,
# In mautrix, send_message returns the event ID string directly
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 +257,111 @@ 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. Initialize SQLite backend for stores
db = Database.create(DB_URI)
await db.start()
# 2. Initialize client
client = Client(base_url=MATRIX_HOMESERVER, mxid=MATRIX_BOT_USER)
# 1. Attempt to restore an existing session
# 3. Session Persistence
session_restored = False
if os.path.exists(CREDENTIALS_FILE):
print("Restoring existing Matrix session...")
log.info("Restoring existing Matrix session...")
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.api.token = creds["access_token"]
client.mxid = 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
log.warning(f"Corrupted credentials file ({e}). Forcing new login...")
if not session_restored:
print(f"No valid session found. Logging in to {MATRIX_HOMESERVER}...")
resp = await client.login(MATRIX_BOT_PASSWORD)
log.info(f"Logging in to {MATRIX_HOMESERVER}...")
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)
with open(CREDENTIALS_FILE, "w") as f:
json.dump({
"access_token": resp.access_token,
"user_id": resp.user_id,
"device_id": resp.device_id,
"access_token": client.api.token,
"user_id": client.mxid,
"device_id": client.device_id,
}, f)
print(f"Logged in and saved credentials for {MATRIX_BOT_USER}.")
log.info("Logged in and saved credentials.")
print("Syncing...")
await client.sync_forever(timeout=30000, full_state=True)
# 4. Attach Crypto & State Stores
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__":

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
markdown