|
-
-
- """
- wechaty channel
- Python Wechaty - https://github.com/wechaty/python-wechaty
- """
- import io
- import os
- import json
- import time
- import asyncio
- import requests
- import pysilk
- import wave
- from pydub import AudioSegment
- from typing import Optional, Union
- from bridge.context import Context, ContextType
- from wechaty_puppet import MessageType, FileBox, ScanStatus
- from wechaty import Wechaty, Contact
- from wechaty.user import Message, Room, MiniProgram, UrlLink
- from channel.channel import Channel
- from common.log import logger
- from common.tmp_dir import TmpDir
- from config import conf
-
-
- class WechatyChannel(Channel):
-
- def __init__(self):
- pass
-
- def startup(self):
- asyncio.run(self.main())
-
- async def main(self):
- config = conf()
-
- token = config.get('wechaty_puppet_service_token')
- os.environ['WECHATY_PUPPET_SERVICE_TOKEN'] = token
- global bot
- bot = Wechaty()
-
- bot.on('scan', self.on_scan)
- bot.on('login', self.on_login)
- bot.on('message', self.on_message)
- await bot.start()
-
- async def on_login(self, contact: Contact):
- logger.info('[WX] login user={}'.format(contact))
-
- async def on_scan(self, status: ScanStatus, qr_code: Optional[str] = None,
- data: Optional[str] = None):
- contact = self.Contact.load(self.contact_id)
- logger.info('[WX] scan user={}, scan status={}, scan qr_code={}'.format(contact, status.name, qr_code))
-
-
- async def on_message(self, msg: Message):
- """
- listen for message event
- """
- from_contact = msg.talker()
- to_contact = msg.to()
- room = msg.room()
- from_user_id = from_contact.contact_id
- to_user_id = to_contact.contact_id
-
- content = msg.text()
- mention_content = await msg.mention_text()
- match_prefix = self.check_prefix(content, conf().get('single_chat_prefix'))
- conversation: Union[Room, Contact] = from_contact if room is None else room
-
- if room is None and msg.type() == MessageType.MESSAGE_TYPE_TEXT:
- if not msg.is_self() and match_prefix is not None:
-
- if match_prefix != '':
- str_list = content.split(match_prefix, 1)
- if len(str_list) == 2:
- content = str_list[1].strip()
-
- img_match_prefix = self.check_prefix(content, conf().get('image_create_prefix'))
- if img_match_prefix:
- content = content.split(img_match_prefix, 1)[1].strip()
- await self._do_send_img(content, from_user_id)
- else:
- await self._do_send(content, from_user_id)
- elif msg.is_self() and match_prefix:
-
- str_list = content.split(match_prefix, 1)
- if len(str_list) == 2:
- content = str_list[1].strip()
- img_match_prefix = self.check_prefix(content, conf().get('image_create_prefix'))
- if img_match_prefix:
- content = content.split(img_match_prefix, 1)[1].strip()
- await self._do_send_img(content, to_user_id)
- else:
- await self._do_send(content, to_user_id)
- elif room is None and msg.type() == MessageType.MESSAGE_TYPE_AUDIO:
- if not msg.is_self():
-
- voice_file = await msg.to_file_box()
- silk_file = TmpDir().path() + voice_file.name
- await voice_file.to_file(silk_file)
- logger.info("[WX]receive voice file: " + silk_file)
-
- wav_file = silk_file.replace(".slk", ".wav")
- with open(silk_file, 'rb') as f:
- silk_data = f.read()
- pcm_data = pysilk.decode(silk_data)
-
- with wave.open(wav_file, 'wb') as wav_data:
- wav_data.setnchannels(1)
- wav_data.setsampwidth(2)
- wav_data.setframerate(24000)
- wav_data.writeframes(pcm_data)
- if os.path.exists(wav_file):
- converter_state = "true"
- else:
- converter_state = "false"
- logger.info("[WX]receive voice converter: " + converter_state)
-
- query = super().build_voice_to_text(wav_file).content
-
- match_prefix = self.check_prefix(query, conf().get('single_chat_prefix'))
- if match_prefix is not None:
- if match_prefix != '':
- str_list = query.split(match_prefix, 1)
- if len(str_list) == 2:
- query = str_list[1].strip()
-
- if conf().get('voice_reply_voice'):
- await self._do_send_voice(query, from_user_id)
- else:
- await self._do_send(query, from_user_id)
- else:
- logger.info("[WX]receive voice check prefix: " + 'False')
-
- os.remove(wav_file)
- os.remove(silk_file)
- elif room and msg.type() == MessageType.MESSAGE_TYPE_TEXT:
-
- room_id = room.room_id
- room_name = await room.topic()
- from_user_id = from_contact.contact_id
- from_user_name = from_contact.name
- is_at = await msg.mention_self()
- content = mention_content
- config = conf()
- match_prefix = (is_at and not config.get("group_at_off", False)) \
- or self.check_prefix(content, config.get('group_chat_prefix')) \
- or self.check_contain(content, config.get('group_chat_keyword'))
-
-
- prefixes = config.get('group_chat_prefix')
- for prefix in prefixes:
- if content.startswith(prefix):
- content = content.replace(prefix, '', 1).strip()
- break
- if ('ALL_GROUP' in config.get('group_name_white_list') or room_name in config.get(
- 'group_name_white_list') or self.check_contain(room_name, config.get(
- 'group_name_keyword_white_list'))) and match_prefix:
- img_match_prefix = self.check_prefix(content, conf().get('image_create_prefix'))
- if img_match_prefix:
- content = content.split(img_match_prefix, 1)[1].strip()
- await self._do_send_group_img(content, room_id)
- else:
- await self._do_send_group(content, room_id, room_name, from_user_id, from_user_name)
- elif room and msg.type() == MessageType.MESSAGE_TYPE_AUDIO:
-
- room_id = room.room_id
- room_name = await room.topic()
- from_user_id = from_contact.contact_id
- from_user_name = from_contact.name
- is_at = await msg.mention_self()
- config = conf()
-
- if config.get('group_speech_recognition') and (
- 'ALL_GROUP' in config.get('group_name_white_list') or room_name in config.get(
- 'group_name_white_list') or self.check_contain(room_name, config.get(
- 'group_name_keyword_white_list'))):
-
- voice_file = await msg.to_file_box()
- silk_file = TmpDir().path() + voice_file.name
- await voice_file.to_file(silk_file)
- logger.info("[WX]receive voice file: " + silk_file)
-
- wav_file = silk_file.replace(".slk", ".wav")
- with open(silk_file, 'rb') as f:
- silk_data = f.read()
- pcm_data = pysilk.decode(silk_data)
-
- with wave.open(wav_file, 'wb') as wav_data:
- wav_data.setnchannels(1)
- wav_data.setsampwidth(2)
- wav_data.setframerate(24000)
- wav_data.writeframes(pcm_data)
- if os.path.exists(wav_file):
- converter_state = "true"
- else:
- converter_state = "false"
- logger.info("[WX]receive voice converter: " + converter_state)
-
- query = super().build_voice_to_text(wav_file).content
-
- match_prefix = self.check_prefix(query, config.get('group_chat_prefix')) \
- or self.check_contain(query, config.get('group_chat_keyword'))
-
- if match_prefix is not None:
-
- prefixes = config.get('group_chat_prefix')
- for prefix in prefixes:
- if query.startswith(prefix):
- query = query.replace(prefix, '', 1).strip()
- break
-
- img_match_prefix = self.check_prefix(query, conf().get('image_create_prefix'))
- if img_match_prefix:
- query = query.split(img_match_prefix, 1)[1].strip()
- await self._do_send_group_img(query, room_id)
- elif config.get('voice_reply_voice'):
- await self._do_send_group_voice(query, room_id, room_name, from_user_id, from_user_name)
- else:
- await self._do_send_group(query, room_id, room_name, from_user_id, from_user_name)
- else:
- logger.info("[WX]receive voice check prefix: " + 'False')
-
- os.remove(wav_file)
- os.remove(silk_file)
-
- async def send(self, message: Union[str, Message, FileBox, Contact, UrlLink, MiniProgram], receiver):
- logger.info('[WX] sendMsg={}, receiver={}'.format(message, receiver))
- if receiver:
- contact = await bot.Contact.find(receiver)
- await contact.say(message)
-
- async def send_group(self, message: Union[str, Message, FileBox, Contact, UrlLink, MiniProgram], receiver):
- logger.info('[WX] sendMsg={}, receiver={}'.format(message, receiver))
- if receiver:
- room = await bot.Room.find(receiver)
- await room.say(message)
-
- async def _do_send(self, query, reply_user_id):
- try:
- if not query:
- return
- context = Context(ContextType.TEXT, query)
- context['session_id'] = reply_user_id
- reply_text = super().build_reply_content(query, context).content
- if reply_text:
- await self.send(conf().get("single_chat_reply_prefix") + reply_text, reply_user_id)
- except Exception as e:
- logger.exception(e)
-
- async def _do_send_voice(self, query, reply_user_id):
- try:
- if not query:
- return
- context = Context(ContextType.TEXT, query)
- context['session_id'] = reply_user_id
- reply_text = super().build_reply_content(query, context).content
- if reply_text:
-
- mp3_file = super().build_text_to_voice(reply_text).content
- silk_file = mp3_file.replace(".mp3", ".silk")
-
- audio = AudioSegment.from_file(mp3_file, format="mp3")
-
- audio = audio.set_frame_rate(24000).set_channels(1)
- wav_data = audio.raw_data
- sample_width = audio.sample_width
-
- silk_data = pysilk.encode(wav_data, 24000)
-
- with open(silk_file, "wb") as f:
- f.write(silk_data)
-
- t = int(time.time())
- file_box = FileBox.from_file(silk_file, name=str(t) + '.silk')
- await self.send(file_box, reply_user_id)
-
- os.remove(mp3_file)
- os.remove(silk_file)
- except Exception as e:
- logger.exception(e)
-
- async def _do_send_img(self, query, reply_user_id):
- try:
- if not query:
- return
- context = Context(ContextType.IMAGE_CREATE, query)
- img_url = super().build_reply_content(query, context).content
- if not img_url:
- return
-
-
-
-
-
-
-
-
- logger.info('[WX] sendImage, receiver={}'.format(reply_user_id))
- t = int(time.time())
- file_box = FileBox.from_url(url=img_url, name=str(t) + '.png')
- await self.send(file_box, reply_user_id)
- except Exception as e:
- logger.exception(e)
-
- async def _do_send_group(self, query, group_id, group_name, group_user_id, group_user_name):
- if not query:
- return
- context = Context(ContextType.TEXT, query)
- group_chat_in_one_session = conf().get('group_chat_in_one_session', [])
- if ('ALL_GROUP' in group_chat_in_one_session or \
- group_name in group_chat_in_one_session or \
- self.check_contain(group_name, group_chat_in_one_session)):
- context['session_id'] = str(group_id)
- else:
- context['session_id'] = str(group_id) + '-' + str(group_user_id)
- reply_text = super().build_reply_content(query, context).content
- if reply_text:
- reply_text = '@' + group_user_name + ' ' + reply_text.strip()
- await self.send_group(conf().get("group_chat_reply_prefix", "") + reply_text, group_id)
-
- async def _do_send_group_voice(self, query, group_id, group_name, group_user_id, group_user_name):
- if not query:
- return
- context = Context(ContextType.TEXT, query)
- group_chat_in_one_session = conf().get('group_chat_in_one_session', [])
- if ('ALL_GROUP' in group_chat_in_one_session or \
- group_name in group_chat_in_one_session or \
- self.check_contain(group_name, group_chat_in_one_session)):
- context['session_id'] = str(group_id)
- else:
- context['session_id'] = str(group_id) + '-' + str(group_user_id)
- reply_text = super().build_reply_content(query, context).content
- if reply_text:
- reply_text = '@' + group_user_name + ' ' + reply_text.strip()
-
- mp3_file = super().build_text_to_voice(reply_text).content
- silk_file = mp3_file.replace(".mp3", ".silk")
-
- audio = AudioSegment.from_file(mp3_file, format="mp3")
-
- audio = audio.set_frame_rate(24000).set_channels(1)
- wav_data = audio.raw_data
- sample_width = audio.sample_width
-
- silk_data = pysilk.encode(wav_data, 24000)
-
- with open(silk_file, "wb") as f:
- f.write(silk_data)
-
- t = int(time.time())
- file_box = FileBox.from_file(silk_file, name=str(t) + '.silk')
- await self.send_group(file_box, group_id)
-
- os.remove(mp3_file)
- os.remove(silk_file)
-
- async def _do_send_group_img(self, query, reply_room_id):
- try:
- if not query:
- return
- context = Context(ContextType.IMAGE_CREATE, query)
- img_url = super().build_reply_content(query, context).content
- if not img_url:
- return
-
- logger.info('[WX] sendImage, receiver={}'.format(reply_room_id))
- t = int(time.time())
- file_box = FileBox.from_url(url=img_url, name=str(t) + '.png')
- await self.send_group(file_box, reply_room_id)
- except Exception as e:
- logger.exception(e)
-
- def check_prefix(self, content, prefix_list):
- for prefix in prefix_list:
- if content.startswith(prefix):
- return prefix
- return None
-
- def check_contain(self, content, keyword_list):
- if not keyword_list:
- return None
- for ky in keyword_list:
- if content.find(ky) != -1:
- return True
- return None
|