cinny/scripts/dev-seed.py

#!/usr/bin/env python3
"""Seed the local dev homeserver (scripts/dev-homeserver.sh) with two users,
a busy unencrypted room and a voice room: alice + bob, "Busy Room" (N
messages, one image every 10th), "Voice Lounge" (org.matrix.msc3417.call with
the call-member power level Lotus applies). Makes alice a server admin (the
voice-limit guard reads room state through the admin API) and writes her
token to .dev-homeserver/admin.token for `dev-homeserver.sh calls`.
Idempotent for the users; every run creates new rooms.

    python3 scripts/dev-seed.py [N=400]
"""
import json, os, sqlite3, struct, sys, time, urllib.error, urllib.parse, urllib.request, zlib

HS = "http://127.0.0.1:8008"
PASSWORD = "password123"


def req(method, path, data=None, token=None, raw=None, ctype="application/json"):
    headers = {"Content-Type": ctype}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    body = raw if raw is not None else (json.dumps(data).encode() if data is not None else None)
    r = urllib.request.Request(HS + path, data=body, headers=headers, method=method)
    return json.load(urllib.request.urlopen(r))


def register_or_login(user):
    try:
        return req("POST", "/_matrix/client/v3/register",
                   {"username": user, "password": PASSWORD, "auth": {"type": "m.login.dummy"}})
    except urllib.error.HTTPError:
        return req("POST", "/_matrix/client/v3/login",
                   {"type": "m.login.password", "identifier": {"type": "m.id.user", "user": user}, "password": PASSWORD})


def png(w, h, rgb):
    raw = b"".join(b"\x00" + bytes(rgb) * w for _ in range(h))
    def chunk(t, d):
        return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xFFFFFFFF)
    return (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
            + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b""))


def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 400
    alice, bob = register_or_login("alice"), register_or_login("bob")
    ta, tb = alice["access_token"], bob["access_token"]
    room = req("POST", "/_matrix/client/v3/createRoom",
               {"name": "Busy Room", "preset": "public_chat", "visibility": "public"}, ta)["room_id"]
    req("POST", f"/_matrix/client/v3/join/{urllib.parse.quote(room)}", {}, tb)
    txn = int(time.time() * 1000)
    for i in range(n):
        tok = ta if i % 2 == 0 else tb
        if i % 10 == 0:
            data = png(64, 48, ((i * 37) % 256, (i * 91) % 256, (i * 17) % 256))
            up = req("POST", f"/_matrix/media/v3/upload?filename=img{i}.png", token=tok, raw=data, ctype="image/png")
            content = {"msgtype": "m.image", "body": f"img{i}.png", "url": up["content_uri"],
                       "info": {"mimetype": "image/png", "w": 64, "h": 48, "size": len(data)}}
        else:
            content = {"msgtype": "m.text", "body": f"message #{i}"}
        txn += 1
        req("PUT", f"/_matrix/client/v3/rooms/{urllib.parse.quote(room)}/send/m.room.message/{txn}", content, tok)
    voice = req("POST", "/_matrix/client/v3/createRoom", {
        "name": "Voice Lounge", "preset": "public_chat",
        "creation_content": {"type": "org.matrix.msc3417.call"},
        "initial_state": [{"type": "org.matrix.msc3401.call", "state_key": "", "content": {}}],
        "power_level_content_override": {"events": {"org.matrix.msc3401.call.member": 0}},
    }, ta)["room_id"]
    req("POST", f"/_matrix/client/v3/join/{urllib.parse.quote(voice)}", {}, tb)
    # admin flag for the guard (takes effect for tokens issued after a Synapse restart-free
    # cache miss; the guard only needs it for the admin state API)
    here = os.path.dirname(os.path.abspath(__file__))
    db = os.path.join(here, "..", ".dev-homeserver", "homeserver.db")
    if os.path.exists(db):
        c = sqlite3.connect(db); c.execute("UPDATE users SET admin=1 WHERE name=?", (alice["user_id"],)); c.commit(); c.close()
        with open(os.path.join(here, "..", ".dev-homeserver", "admin.token"), "w") as f:
            f.write(ta)
    print(json.dumps({"room": room, "voice_room": voice, "alice": alice["user_id"], "bob": bob["user_id"], "count": n}))


if __name__ == "__main__":
    main()