mirror of
https://github.com/xfgryujk/blivechat.git
synced 2024-12-26 21:00:15 +08:00
279 lines
7.3 KiB
Python
279 lines
7.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
import dataclasses
|
|
import datetime
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import string
|
|
import subprocess
|
|
from typing import *
|
|
|
|
import api.plugin
|
|
import blcsdk
|
|
import blcsdk.models as sdk_models
|
|
import config
|
|
import update
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
PLUGINS_PATH = os.path.join(config.DATA_PATH, 'plugins')
|
|
|
|
_plugins: Dict[str, 'Plugin'] = {}
|
|
|
|
|
|
def init():
|
|
plugin_ids = _discover_plugin_ids()
|
|
if not plugin_ids:
|
|
return
|
|
logger.info('Found plugins: %s', plugin_ids)
|
|
|
|
for plugin_id in plugin_ids:
|
|
plugin = _create_plugin(plugin_id)
|
|
if plugin is not None:
|
|
_plugins[plugin_id] = plugin
|
|
|
|
for plugin in _plugins.values():
|
|
if plugin.enabled:
|
|
try:
|
|
plugin.start()
|
|
except SwitchPluginError:
|
|
pass
|
|
|
|
|
|
def shut_down():
|
|
for plugin in _plugins.values():
|
|
plugin.stop()
|
|
|
|
|
|
def _discover_plugin_ids():
|
|
res = []
|
|
try:
|
|
with os.scandir(PLUGINS_PATH) as it:
|
|
for entry in it:
|
|
if entry.is_dir() and os.path.isfile(os.path.join(entry.path, 'plugin.json')):
|
|
res.append(entry.name)
|
|
except OSError:
|
|
logger.exception('Failed to discover plugins:')
|
|
return res
|
|
|
|
|
|
def _create_plugin(plugin_id):
|
|
config_path = os.path.join(PLUGINS_PATH, plugin_id, 'plugin.json')
|
|
try:
|
|
plugin_config = PluginConfig.from_file(config_path)
|
|
except (OSError, json.JSONDecodeError, TypeError):
|
|
logger.exception('plugin=%s failed to load config:', plugin_id)
|
|
return None
|
|
return Plugin(plugin_id, plugin_config)
|
|
|
|
|
|
def iter_plugins() -> Iterable['Plugin']:
|
|
return _plugins.values()
|
|
|
|
|
|
def get_plugin(plugin_id):
|
|
return _plugins.get(plugin_id, None)
|
|
|
|
|
|
def get_plugin_by_token(token):
|
|
if token == '':
|
|
return None
|
|
# 应该最多就十几个插件吧，偷懒用遍历了
|
|
for plugin in _plugins.values():
|
|
if plugin.token == token:
|
|
return plugin
|
|
return None
|
|
|
|
|
|
def broadcast_cmd_data(cmd, data, extra: Optional[dict] = None):
|
|
body = api.plugin.make_message_body(cmd, data, extra)
|
|
for plugin in _plugins.values():
|
|
plugin.send_body_no_raise(body)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class PluginConfig:
|
|
name: str = ''
|
|
version: str = ''
|
|
author: str = ''
|
|
description: str = ''
|
|
run_cmd: str = ''
|
|
enabled: bool = False
|
|
|
|
@classmethod
|
|
def from_file(cls, path):
|
|
with open(path, encoding='utf-8') as f:
|
|
cfg = json.load(f)
|
|
if not isinstance(cfg, dict):
|
|
raise TypeError(f'Config type error, type={type(cfg)}')
|
|
|
|
return cls(
|
|
name=str(cfg.get('name', '')),
|
|
version=str(cfg.get('version', '')),
|
|
author=str(cfg.get('author', '')),
|
|
description=str(cfg.get('description', '')),
|
|
run_cmd=str(cfg.get('run', '')),
|
|
enabled=bool(cfg.get('enabled', False)),
|
|
)
|
|
|
|
def save(self, path):
|
|
try:
|
|
with open(path, encoding='utf-8') as f:
|
|
cfg = json.load(f)
|
|
if not isinstance(cfg, dict):
|
|
raise TypeError(f'Config type error, type={type(cfg)}')
|
|
except (OSError, json.JSONDecodeError, TypeError):
|
|
cfg = {}
|
|
|
|
cfg['name'] = self.name
|
|
cfg['version'] = self.version
|
|
cfg['author'] = self.author
|
|
cfg['description'] = self.description
|
|
cfg['run'] = self.run_cmd
|
|
cfg['enabled'] = self.enabled
|
|
|
|
tmp_path = path + '.tmp'
|
|
with open(tmp_path, 'w', encoding='utf-8') as f:
|
|
json.dump(cfg, f, ensure_ascii=False, indent=2)
|
|
os.replace(tmp_path, path)
|
|
|
|
|
|
class SwitchPluginError(Exception):
|
|
"""开关插件时错误"""
|
|
|
|
|
|
class SwitchTooFrequently(SwitchPluginError):
|
|
"""开关插件太频繁"""
|
|
|
|
|
|
class Plugin:
|
|
def __init__(self, plugin_id, cfg: PluginConfig):
|
|
self._id = plugin_id
|
|
self._config = cfg
|
|
|
|
self._last_switch_time = datetime.datetime.fromtimestamp(0)
|
|
self._token = ''
|
|
self._client: Optional['api.plugin.PluginWsHandler'] = None
|
|
|
|
@property
|
|
def id(self):
|
|
return self._id
|
|
|
|
@property
|
|
def config(self):
|
|
return self._config
|
|
|
|
@property
|
|
def enabled(self):
|
|
return self._config.enabled
|
|
|
|
@enabled.setter
|
|
def enabled(self, value):
|
|
if self._config.enabled == value:
|
|
return
|
|
self._config.enabled = value
|
|
|
|
config_path = os.path.join(self.base_path, 'plugin.json')
|
|
try:
|
|
self._config.save(config_path)
|
|
except OSError:
|
|
logger.exception('plugin=%s failed to save config', self._id)
|
|
|
|
if value:
|
|
self.start()
|
|
else:
|
|
self.stop()
|
|
|
|
@property
|
|
def base_path(self):
|
|
return os.path.join(PLUGINS_PATH, self._id)
|
|
|
|
@property
|
|
def token(self):
|
|
return self._token
|
|
|
|
@property
|
|
def is_started(self):
|
|
return self._token != ''
|
|
|
|
@property
|
|
def is_connected(self):
|
|
return self._client is not None
|
|
|
|
def start(self):
|
|
if self.is_started:
|
|
return
|
|
self._refresh_last_switch_time()
|
|
|
|
token = ''.join(random.choice(string.hexdigits) for _ in range(32))
|
|
self._set_token(token)
|
|
|
|
cfg = config.get_config()
|
|
env = {
|
|
**os.environ,
|
|
'BLC_PORT': str(cfg.port),
|
|
'BLC_TOKEN': self._token,
|
|
}
|
|
try:
|
|
subprocess.Popen(
|
|
self._config.run_cmd,
|
|
shell=True,
|
|
cwd=self.base_path,
|
|
env=env,
|
|
)
|
|
except OSError as e:
|
|
logger.exception('plugin=%s failed to start', self._id)
|
|
raise SwitchPluginError(str(e))
|
|
|
|
def _refresh_last_switch_time(self):
|
|
cur_time = datetime.datetime.now()
|
|
if cur_time - self._last_switch_time < datetime.timedelta(seconds=3):
|
|
raise SwitchTooFrequently(f'plugin={self._id} switches too frequently')
|
|
self._last_switch_time = cur_time
|
|
|
|
def stop(self):
|
|
if not self.is_started:
|
|
return
|
|
self._refresh_last_switch_time()
|
|
|
|
self._set_token('')
|
|
|
|
def _set_token(self, token):
|
|
if self._token == token:
|
|
return
|
|
self._token = token
|
|
|
|
# 踢掉已经连接的客户端
|
|
self._set_client(None)
|
|
|
|
def _set_client(self, client: Optional['api.plugin.PluginWsHandler']):
|
|
if self._client is client:
|
|
return
|
|
if self._client is not None:
|
|
logger.info('plugin=%s closing old client', self._id)
|
|
self._client.close()
|
|
self._client = client
|
|
|
|
def on_client_connect(self, client: 'api.plugin.PluginWsHandler'):
|
|
self._set_client(client)
|
|
|
|
# 发送初始化消息
|
|
self.send_cmd_data(sdk_models.Command.BLC_INIT, {
|
|
'blcVersion': update.VERSION,
|
|
'sdkVersion': blcsdk.__version__,
|
|
'pluginId': self._id,
|
|
})
|
|
|
|
def on_client_close(self, client: 'api.plugin.PluginWsHandler'):
|
|
if self._client is client:
|
|
self._set_client(None)
|
|
|
|
def send_cmd_data(self, cmd, data, extra: Optional[dict] = None):
|
|
if self._client is not None:
|
|
self._client.send_cmd_data(cmd, data, extra)
|
|
|
|
def send_body_no_raise(self, body):
|
|
if self._client is not None:
|
|
self._client.send_body_no_raise(body)
|