mirror of
https://github.com/famedly/ansible-collection-matrix
synced 2025-01-10 00:28:42 +00:00
115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
import json
|
|
import secrets
|
|
import string
|
|
from copy import deepcopy
|
|
from typing import Any, Optional, Union
|
|
from typing import Dict
|
|
|
|
|
|
class RoomEvents:
|
|
M_ROOM_CREATE = {
|
|
"content": {
|
|
"creator": "@example:localhost",
|
|
},
|
|
"sender": "@example:localhost",
|
|
"state_key": "",
|
|
"type": "m.room.create",
|
|
}
|
|
M_ROOM_MEMBER = {
|
|
"content": {"membership": "join"},
|
|
"membership": "join",
|
|
"sender": "@example:localhost",
|
|
"state_key": "@example:localhost",
|
|
"type": "m.room.member",
|
|
}
|
|
# This is not to spec!
|
|
M_ROOM_DUMMY = {
|
|
"content": {"dummy": "I'm a test message"},
|
|
"sender": "@example:localhost",
|
|
"state_key": "@example:localhost",
|
|
"type": "m.room.dummy",
|
|
}
|
|
|
|
|
|
class RoomSimulator:
|
|
INITIAL_DATA = {"rooms": {}, "room_directory": {}}
|
|
|
|
def __init__(self, data: str = json.dumps(INITIAL_DATA)):
|
|
try:
|
|
self.data = json.loads(data)
|
|
except (json.JSONDecodeError, TypeError):
|
|
self.data = self.INITIAL_DATA
|
|
|
|
def random_event_id(self):
|
|
return "".join(
|
|
secrets.choice(string.ascii_uppercase + string.digits) for i in range(64)
|
|
)
|
|
|
|
def send_event(self, room_id: str, event: Dict[str, Any]) -> Dict[str, Any]:
|
|
event["event_id"] = self.random_event_id()
|
|
for e in self.get_events(room_id):
|
|
if e["type"] == event["type"] and e["state_key"] == event["state_key"]:
|
|
event["replaces_state"] = e["event_id"]
|
|
self.add_event(room_id=room_id, event=event)
|
|
return event
|
|
|
|
def get_events(self, room_id: str) -> Dict[str, Any]:
|
|
return self.data["rooms"][room_id]
|
|
|
|
def add_event(self, room_id: str, event: Dict[str, Any]) -> Dict[str, Any]:
|
|
return self.data["rooms"][room_id].append(event)
|
|
|
|
def get_state_event(
|
|
self, room_id: str, event_type: str, state_key: str
|
|
) -> Union[None, Dict[str, Any]]:
|
|
events = list[Dict[str, Any]]()
|
|
for e in self.get_events(room_id):
|
|
if e["type"] == event_type and e["state_key"] == state_key:
|
|
events.append(e)
|
|
if len(events) > 0:
|
|
return events[-1]
|
|
return None
|
|
|
|
def add_room(self, room_id: str, alias: Optional[str] = None):
|
|
self.data["rooms"][room_id] = []
|
|
if alias is not None:
|
|
self.data["room_directory"][alias] = room_id
|
|
return room_id
|
|
|
|
def resolve_alias(self, alias: str) -> str:
|
|
return self.data["room_directory"].get(alias, None)
|
|
|
|
def export(self) -> str:
|
|
return json.dumps(self.data)
|
|
|
|
# Matrix Tasks
|
|
|
|
def create_room(self, alias: Optional[str] = None):
|
|
room_id = (
|
|
"!"
|
|
+ "".join(
|
|
secrets.choice(string.ascii_uppercase + string.digits)
|
|
for i in range(10)
|
|
)
|
|
+ ":matrix.example.tld"
|
|
)
|
|
self.add_room(room_id, alias)
|
|
return room_id
|
|
|
|
# Matrix Events
|
|
|
|
def m_room_create(self, room_id: str, creator: str) -> Dict[str, Any]:
|
|
data = deepcopy(RoomEvents.M_ROOM_CREATE)
|
|
data["content"]["creator"] = creator
|
|
data["sender"] = creator
|
|
return self.send_event(room_id=room_id, event=data)
|
|
|
|
def m_room_member(
|
|
self, room_id: str, membership: str, sender: str, state_key: str
|
|
) -> Dict[str, Any]:
|
|
data = deepcopy(RoomEvents.M_ROOM_MEMBER)
|
|
data["content"]["membership"] = membership
|
|
data["membership"] = membership
|
|
data["sender"] = sender
|
|
data["state_key"] = state_key
|
|
return self.send_event(room_id=room_id, event=data)
|