blivedm/sample.py

import asyncio
import http.cookies
import json
import logging
import os
import signal
import threading
import time
from typing import *

import aiohttp
import redis

import blivedm
from blivedm.clients import ws_base

with open("config.json") as conf_f:
    conf = json.load(conf_f)

    redis_conf = conf.get("redis", {})

os.makedirs("logs", exist_ok=True)
room_status_log = open("logs/room_status.jsonl", "a", encoding="utf-8")

# set log level to INFO
logging.basicConfig(level=logging.INFO)

logger = logging.getLogger()

room_list_key = "bilibili.live.danmu.room_list"
r = redis.Redis(
    host=redis_conf.get("host", "127.0.0.1"),
    port=redis_conf.get("port", 6379),
    db=redis_conf.get("port", 0),
)

TEST_ROOM_IDS = [int(room_id) for room_id in set(r.lrange(room_list_key, 0, -1))]
clients = {}
log_files = {}

with open("cookie.txt", "r", encoding="utf-8") as f:
    SESSDATA = f.read().strip()

session: Optional[aiohttp.ClientSession] = None


def handle_sighup(signum, frame):
    global log_files
    global room_status_log

    for room_id in log_files:
        try:
            log_files[room_id].close()
        except Exception as e:
            logger.exception(f"close log_files[{room_id}] failed with {e}")
    try:
        room_status_log.close()
    except Exception as e:
        logger.exception(f"close room_status_log failed with {e}")

    log_files = {}
    room_status_log = open("logs/room_status.jsonl", "a", encoding="utf-8")


signal.signal(signal.SIGHUP, handle_sighup)


async def main():
    flush_files()

    init_session()
    try:
        for room_id in TEST_ROOM_IDS:
            client = None
            try:
                client = await connect_room(room_id)
            except Exception as e:
                if client is not None:
                    await client.stop_and_close()

                raise e

        # run subscribe_redis on new thread
        threading.Thread(target=subscribe_redis, daemon=True).start()

        while True:
            await asyncio.gather(*(
                clients[room_id].join() for room_id in clients
            ))
            await asyncio.sleep(1)
    finally:
        for room_id in clients:
            await clients[room_id].stop_and_close()
        await session.close()


async def connect_room(room_id):
    logger.info("connect to room %s", room_id)
    client = blivedm.BLiveClient(room_id, session=session)
    client.set_handler(jsonl_logger)
    await client.init_room()
    client._need_init_room = False
    client.start()
    clients[room_id] = client
    logger.info("connect to room %s success", room_id)

    try:
        room_status_log.write(json.dumps(
            {
                "room_id": room_id,
                "live_status": client.live_status,
                "live_start_time": client.live_start_time,
            },
            separators=(",", ":"),
            ensure_ascii=False,
        ))
        room_status_log.write("\n")
    except Exception as e:
        logger.exception(f"write room_status_log failed with {e}")
    return client


def flush_files():
    try:
        room_status_log.flush()
        for room_id in log_files:
            log_files[room_id].flush()
    except Exception as e:
        logger.exception(f"flush_files failed with {e}")

    timer = threading.Timer(1, flush_files)
    timer.daemon = True
    timer.start()


def subscribe_redis():
    p = r.pubsub()
    p.psubscribe(**{f"__keyspace@0__:{room_list_key}": room_changed_handler})
    for message in p.listen():
        pass


# noinspection PyUnusedLocal
def room_changed_handler(message):
    global TEST_ROOM_IDS
    TEST_ROOM_IDS = set(int(room_id) for room_id in set(r.lrange(room_list_key, 0, -1)))

    for removed in clients.keys() - TEST_ROOM_IDS:
        asyncio.run(clients[removed].stop_and_close())
        del clients[removed]

        if removed in log_files:
            log_files[removed].close()
            del log_files[removed]

    for added in TEST_ROOM_IDS - clients.keys():
        client = None
        try:
            client = asyncio.run(connect_room(added))
        except Exception as e:
            if client is not None:
                asyncio.run(client.stop_and_close())


def init_session():
    cookies = http.cookies.SimpleCookie()
    cookies["SESSDATA"] = SESSDATA
    cookies["SESSDATA"]["domain"] = "bilibili.com"

    global session
    session = aiohttp.ClientSession()
    session.cookie_jar.update_cookies(cookies)


class JsonlLogger(blivedm.BaseHandler):
    def _global_callback(self, client: ws_base.WebSocketClientBase, message: dict):
        if client.room_id not in log_files:
            log_files[client.room_id] = open(f"logs/{client.room_id}.jsonl", "a", encoding="utf-8")

        if "cmd" in message and message["cmd"] == "PREPARING":
            message["timestamp"] = int(time.time())

        log_files[client.room_id].write(f"{json.dumps(message, separators=(',', ':'), ensure_ascii=False)}\n")


jsonl_logger = JsonlLogger()

if __name__ == "__main__":
    asyncio.run(main())