462 lines
16 KiB
Python
462 lines
16 KiB
Python
import asyncio
|
||
import os
|
||
import sys
|
||
import markdown
|
||
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,
|
||
)
|
||
|
||
# ------------------------------------------------------------------------------
|
||
# Configuration
|
||
# ------------------------------------------------------------------------------
|
||
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:nest.local")
|
||
ADMIN_MATRIX_IDS = set(os.getenv("ADMIN_MATRIX_IDS", "@you:nest.local").split(","))
|
||
|
||
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")
|
||
|
||
# In-memory tracking of active requests
|
||
pending_requests = {}
|
||
|
||
|
||
# ------------------------------------------------------------------------------
|
||
# Keycloak Helper Functions
|
||
# ------------------------------------------------------------------------------
|
||
def get_keycloak_client() -> KeycloakAdmin:
|
||
"""Initializes KeycloakAdmin client using client credentials."""
|
||
return KeycloakAdmin(
|
||
server_url=KEYCLOAK_URL,
|
||
client_id=KEYCLOAK_CLIENT_ID,
|
||
client_secret_key=KEYCLOAK_CLIENT_SECRET,
|
||
realm_name=KEYCLOAK_REALM,
|
||
user_realm_name=KEYCLOAK_REALM,
|
||
grant_type="client_credentials",
|
||
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})
|
||
if not users:
|
||
return False, f"⚠️ Keycloak user `{username}` not found."
|
||
|
||
user_id = users[0]["id"]
|
||
user_groups = kc.get_user_groups(user_id=user_id)
|
||
user_group_names = {g["name"].lower() for g in user_groups}
|
||
|
||
if group_name.lower() in user_group_names:
|
||
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."
|
||
user_id = users[0]["id"]
|
||
|
||
groups = kc.get_groups()
|
||
target_group = next((g for g in groups if g["name"].lower() == group_name.lower()), None)
|
||
if not target_group:
|
||
return False, f"Keycloak group '{group_name}' not found."
|
||
group_id = target_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']}'."
|
||
|
||
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()
|
||
else:
|
||
user_id = users[0]["id"]
|
||
user_groups = kc.get_user_groups(user_id=user_id)
|
||
user_group_names = {g["name"] for g in user_groups}
|
||
|
||
groups = kc.get_groups()
|
||
valid_groups = [g["name"] for g in groups if g["name"].lower() != "admin"]
|
||
|
||
if not valid_groups:
|
||
return True, "No groups are currently available to request."
|
||
|
||
group_lines = []
|
||
for name in valid_groups:
|
||
if name in user_group_names:
|
||
group_lines.append(f"* `{name}` *(Already a member)*")
|
||
else:
|
||
group_lines.append(f"* `{name}`")
|
||
|
||
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)}"
|
||
|
||
|
||
# ------------------------------------------------------------------------------
|
||
# Auto-Verification Event Handler
|
||
# ------------------------------------------------------------------------------
|
||
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):
|
||
print(f"Key verification MAC received from {event.sender} ({event.from_device}). Finalizing verification...")
|
||
|
||
# 1. Send m.key.verification.done to complete the SAS protocol in Element Web
|
||
done_msg = ToDeviceMessage(
|
||
"m.key.verification.done",
|
||
event.sender,
|
||
event.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 device as verified in matrix-nio's device store
|
||
if event.sender in client.device_store and event.from_device in client.device_store[event.sender]:
|
||
device = client.device_store[event.sender][event.from_device]
|
||
client.verify_device(device)
|
||
print(f"✅ Device {event.from_device} for {event.sender} is now verified and trusted!")
|
||
|
||
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: str = "m.text",
|
||
relates_to: dict = None,
|
||
):
|
||
html_body = markdown.markdown(md_text, extensions=["extra", "sane_lists", "nl2br"])
|
||
|
||
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)
|
||
|
||
|
||
# ------------------------------------------------------------------------------
|
||
# 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
|
||
|
||
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 = 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(
|
||
client,
|
||
room_id=room.room_id,
|
||
md_text=msg,
|
||
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(
|
||
client,
|
||
room_id=room.room_id,
|
||
md_text="Usage: `!request <group_name>` or `!groups`",
|
||
msgtype="m.notice",
|
||
)
|
||
return
|
||
|
||
requested_group = parts[1].strip()
|
||
|
||
if requested_group.lower() == "admin":
|
||
await send_markdown_message(
|
||
client,
|
||
room_id=room.room_id,
|
||
md_text="❌ Requests for the **admin** group are not permitted.",
|
||
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(
|
||
client,
|
||
room_id=room.room_id,
|
||
md_text=membership_msg,
|
||
msgtype="m.notice",
|
||
)
|
||
return
|
||
|
||
admin_msg_body = (
|
||
f"📋 **Access Request**\n\n"
|
||
f"**User:** `{matrix_user}` (KC: `{kc_username}`)\n"
|
||
f"**Requested Group:** `{requested_group}`\n\n"
|
||
f"React with 👍 to **Approve** or 👎 to **Deny**."
|
||
)
|
||
|
||
res = await send_markdown_message(
|
||
client,
|
||
room_id=ADMIN_ROOM_ID,
|
||
md_text=admin_msg_body,
|
||
)
|
||
|
||
if isinstance(res, RoomSendResponse):
|
||
pending_requests[res.event_id] = {
|
||
"matrix_user": matrix_user,
|
||
"kc_username": kc_username,
|
||
"requested_group": requested_group,
|
||
"origin_room_id": room.room_id,
|
||
}
|
||
|
||
await send_markdown_message(
|
||
client,
|
||
room_id=room.room_id,
|
||
md_text=f"Request for group `{requested_group}` submitted for admin approval.",
|
||
msgtype="m.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:
|
||
return
|
||
|
||
target_event_id = event.reacts_to
|
||
emoji = event.key
|
||
|
||
if not target_event_id or target_event_id not in pending_requests:
|
||
return
|
||
|
||
request = pending_requests[target_event_id]
|
||
|
||
if "👍" in emoji:
|
||
success, msg = add_user_to_kc_group(
|
||
request["kc_username"], request["requested_group"]
|
||
)
|
||
|
||
status_text = (
|
||
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="m.notice",
|
||
)
|
||
|
||
elif "👎" in emoji:
|
||
status_text = (
|
||
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="m.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},
|
||
)
|
||
|
||
del pending_requests[target_event_id]
|
||
|
||
|
||
# ------------------------------------------------------------------------------
|
||
# Main Event Loop
|
||
# ------------------------------------------------------------------------------
|
||
async def main():
|
||
client = AsyncClient(
|
||
MATRIX_HOMESERVER,
|
||
MATRIX_BOT_USER,
|
||
store_path="/app/store",
|
||
device_id="KEYCLOAK_ACCESS_BOT",
|
||
)
|
||
|
||
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,
|
||
),
|
||
)
|
||
|
||
print(f"Connecting to Matrix homeserver {MATRIX_HOMESERVER}...")
|
||
await client.login(MATRIX_BOT_PASSWORD)
|
||
print(f"Logged in as {MATRIX_BOT_USER}. Syncing...")
|
||
|
||
await client.sync_forever(timeout=30000, full_state=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
asyncio.run(main())
|
||
except KeyboardInterrupt:
|
||
print("\nBot stopped by user.")
|
||
sys.exit(0)
|