You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

222 satır
8.2KB

  1. # encoding:utf-8
  2. """
  3. wechat channel
  4. """
  5. import itchat
  6. import json
  7. from itchat.content import *
  8. from channel.channel import Channel
  9. from concurrent.futures import ThreadPoolExecutor
  10. from common.log import logger
  11. from common.tmp_dir import TmpDir
  12. from config import conf
  13. import requests
  14. import io
  15. thread_pool = ThreadPoolExecutor(max_workers=8)
  16. @itchat.msg_register(TEXT)
  17. def handler_single_msg(msg):
  18. WechatChannel().handle_text(msg)
  19. return None
  20. @itchat.msg_register(TEXT, isGroupChat=True)
  21. def handler_group_msg(msg):
  22. WechatChannel().handle_group(msg)
  23. return None
  24. @itchat.msg_register(VOICE)
  25. def handler_single_voice(msg):
  26. WechatChannel().handle_voice(msg)
  27. return None
  28. class WechatChannel(Channel):
  29. def __init__(self):
  30. pass
  31. def startup(self):
  32. # login by scan QRCode
  33. itchat.auto_login(enableCmdQR=2)
  34. # start message listener
  35. itchat.run()
  36. def handle_voice(self, msg):
  37. if conf().get('speech_recognition') != True :
  38. return
  39. logger.debug("[WX]receive voice msg: " + msg['FileName'])
  40. thread_pool.submit(self._do_handle_voice, msg)
  41. def _do_handle_voice(self, msg):
  42. from_user_id = msg['FromUserName']
  43. other_user_id = msg['User']['UserName']
  44. if from_user_id == other_user_id:
  45. file_name = TmpDir().path() + msg['FileName']
  46. msg.download(file_name)
  47. query = super().build_voice_to_text(file_name)
  48. if conf().get('voice_reply_voice'):
  49. self._do_send_voice(query, from_user_id)
  50. else:
  51. self._do_send_text(query, from_user_id)
  52. def handle_text(self, msg):
  53. logger.debug("[WX]receive text msg: " + json.dumps(msg, ensure_ascii=False))
  54. content = msg['Text']
  55. self._handle_single_msg(msg, content)
  56. def _handle_single_msg(self, msg, content):
  57. from_user_id = msg['FromUserName']
  58. to_user_id = msg['ToUserName'] # 接收人id
  59. other_user_id = msg['User']['UserName'] # 对手方id
  60. match_prefix = self.check_prefix(content, conf().get('single_chat_prefix'))
  61. if "」\n- - - - - - - - - - - - - - -" in content:
  62. logger.debug("[WX]reference query skipped")
  63. return
  64. if from_user_id == other_user_id and match_prefix is not None:
  65. # 好友向自己发送消息
  66. if match_prefix != '':
  67. str_list = content.split(match_prefix, 1)
  68. if len(str_list) == 2:
  69. content = str_list[1].strip()
  70. img_match_prefix = self.check_prefix(content, conf().get('image_create_prefix'))
  71. if img_match_prefix:
  72. content = content.split(img_match_prefix, 1)[1].strip()
  73. thread_pool.submit(self._do_send_img, content, from_user_id)
  74. else :
  75. thread_pool.submit(self._do_send_text, content, from_user_id)
  76. elif to_user_id == other_user_id and match_prefix:
  77. # 自己给好友发送消息
  78. str_list = content.split(match_prefix, 1)
  79. if len(str_list) == 2:
  80. content = str_list[1].strip()
  81. img_match_prefix = self.check_prefix(content, conf().get('image_create_prefix'))
  82. if img_match_prefix:
  83. content = content.split(img_match_prefix, 1)[1].strip()
  84. thread_pool.submit(self._do_send_img, content, to_user_id)
  85. else:
  86. thread_pool.submit(self._do_send_text, content, to_user_id)
  87. def handle_group(self, msg):
  88. logger.debug("[WX]receive group msg: " + json.dumps(msg, ensure_ascii=False))
  89. group_name = msg['User'].get('NickName', None)
  90. group_id = msg['User'].get('UserName', None)
  91. if not group_name:
  92. return ""
  93. origin_content = msg['Content']
  94. content = msg['Content']
  95. content_list = content.split(' ', 1)
  96. context_special_list = content.split('\u2005', 1)
  97. if len(context_special_list) == 2:
  98. content = context_special_list[1]
  99. elif len(content_list) == 2:
  100. content = content_list[1]
  101. if "」\n- - - - - - - - - - - - - - -" in content:
  102. logger.debug("[WX]reference query skipped")
  103. return ""
  104. config = conf()
  105. match_prefix = (msg['IsAt'] and not config.get("group_at_off", False)) or self.check_prefix(origin_content, config.get('group_chat_prefix')) \
  106. or self.check_contain(origin_content, config.get('group_chat_keyword'))
  107. if ('ALL_GROUP' in config.get('group_name_white_list') or group_name in config.get('group_name_white_list') or self.check_contain(group_name, config.get('group_name_keyword_white_list'))) and match_prefix:
  108. img_match_prefix = self.check_prefix(content, conf().get('image_create_prefix'))
  109. if img_match_prefix:
  110. content = content.split(img_match_prefix, 1)[1].strip()
  111. thread_pool.submit(self._do_send_img, content, group_id)
  112. else:
  113. thread_pool.submit(self._do_send_group, content, msg)
  114. def send(self, msg, receiver):
  115. itchat.send(msg, toUserName=receiver)
  116. logger.info('[WX] sendMsg={}, receiver={}'.format(msg, receiver))
  117. def _do_send_voice(self, query, reply_user_id):
  118. try:
  119. if not query:
  120. return
  121. context = dict()
  122. context['from_user_id'] = reply_user_id
  123. reply_text = super().build_reply_content(query, context)
  124. if reply_text:
  125. replyFile = super().build_text_to_voice(reply_text)
  126. itchat.send_file(replyFile, toUserName=reply_user_id)
  127. logger.info('[WX] sendFile={}, receiver={}'.format(replyFile, reply_user_id))
  128. except Exception as e:
  129. logger.exception(e)
  130. def _do_send_text(self, query, reply_user_id):
  131. try:
  132. if not query:
  133. return
  134. context = dict()
  135. context['session_id'] = reply_user_id
  136. reply_text = super().build_reply_content(query, context)
  137. if reply_text:
  138. self.send(conf().get("single_chat_reply_prefix") + reply_text, reply_user_id)
  139. except Exception as e:
  140. logger.exception(e)
  141. def _do_send_img(self, query, reply_user_id):
  142. try:
  143. if not query:
  144. return
  145. context = dict()
  146. context['type'] = 'IMAGE_CREATE'
  147. img_url = super().build_reply_content(query, context)
  148. if not img_url:
  149. return
  150. # 图片下载
  151. pic_res = requests.get(img_url, stream=True)
  152. image_storage = io.BytesIO()
  153. for block in pic_res.iter_content(1024):
  154. image_storage.write(block)
  155. image_storage.seek(0)
  156. # 图片发送
  157. itchat.send_image(image_storage, reply_user_id)
  158. logger.info('[WX] sendImage, receiver={}'.format(reply_user_id))
  159. except Exception as e:
  160. logger.exception(e)
  161. def _do_send_group(self, query, msg):
  162. if not query:
  163. return
  164. context = dict()
  165. group_name = msg['User']['NickName']
  166. group_id = msg['User']['UserName']
  167. group_chat_in_one_session = conf().get('group_chat_in_one_session', [])
  168. if ('ALL_GROUP' in group_chat_in_one_session or \
  169. group_name in group_chat_in_one_session or \
  170. self.check_contain(group_name, group_chat_in_one_session)):
  171. context['session_id'] = group_id
  172. else:
  173. context['session_id'] = msg['ActualUserName']
  174. reply_text = super().build_reply_content(query, context)
  175. if reply_text:
  176. reply_text = '@' + msg['ActualNickName'] + ' ' + reply_text.strip()
  177. self.send(conf().get("group_chat_reply_prefix", "") + reply_text, group_id)
  178. def check_prefix(self, content, prefix_list):
  179. for prefix in prefix_list:
  180. if content.startswith(prefix):
  181. return prefix
  182. return None
  183. def check_contain(self, content, keyword_list):
  184. if not keyword_list:
  185. return None
  186. for ky in keyword_list:
  187. if content.find(ky) != -1:
  188. return True
  189. return None