BOTNAME = os.getenv("SHAMAN_NAME", "shaman") BOTPASSWORD = os.getenv("SHAMAN_PASSWORD", "shaman") BOTNAME_NOCASE = re.compile(re.escape(BOTNAME), re.IGNORECASE) if SITE_URL == "": print("Set SHAMAN_SITEURL") sys.exit(1) if BOTPASSWORD == "": print("Set SHAMAN_PASSWORD") sys.exit(1) print("Connecting to.. {}".format(SITE_URL)) RC = RocketChatAPI(settings={'username': BOTNAME, 'password': BOTPASSWORD, 'domain': 'https://' + SITE_URL}) ROOMS = {r['name']: r['id'] for r in RC.get_public_rooms()} ws = None # global rocket chat websocket connection object # websockets rocket chat login handshake def wslogin(): # get an auth token from the rest api auth_token = GetMe(RC.settings).auth_token url = "wss://" + SITE_URL + "/websocket" trace = uuid.uuid4().hex[:5] global ws ws = create_connection(url) # must ping first ws.send(json.dumps({ "msg": "connect",
class Bot(): """ 创建一个聊天机器人 :RoomsDictUpdate: 更新聊天室字典 :SetRoom: 指定使用本聊天机器人的聊天室 """ def __init__(self, username, password, domain): # 设置队列,记录最近答复过的消息id,避免重复回答同一条消息,假定1分钟内@bot的人次不超过100 self.replied_queue = deque([], maxlen=100) # 设置rocketchat的api参数 self.api = RocketChatAPI(settings={'username': username, 'password': password, 'domain': domain}) def _RoomsDict(self, _class='private'): """ 获取聊天房间的字典,字典的key为聊天室名,value为聊天室id,聊天室不存在重名 :_class: 聊天室类型 - private: 私有聊天室 - public: 公共聊天室 """ if _class == 'private': rooms = self.api.get_private_rooms() elif _class == 'public': rooms = self.api.get_public_rooms() else: raise ValueError('请指定正确的_class') fun = lambda d:tuple(d.values()) rooms = {fun(d)[0]:fun(d)[1] for d in rooms} return rooms def RoomsDictUpdate(self): private_rooms = self._RoomsDict('private') public_rooms = self._RoomsDict('public') self.rooms_dict = dict(private_rooms, **public_rooms) def SetRoom(self, room_name): try: self.room_id = self.rooms_dict[room_name] except KeyError: raise KeyError('房间名<%s>不存在' % room_name) def History(self, oldest=None): """ 获取当前聊天室的历史聊天记录 """ history = self.api.get_private_room_history(self.room_id, oldest=oldest) if history['success']: history_msgs = history['messages'] return history_msgs def NewMsgs(self): """ 获取指定聊天室最近60秒的历史消息 """ begin_CCT = datetime.now() + timedelta(seconds=-10) begin_UTC = begin_CCT.astimezone(pytz.timezone('UTC')) oldest = begin_UTC.strftime('%Y-%m-%dT%H:%M:%SZ') msgs = self.History(oldest=oldest) msgs.reverse() # 逆序排列 return msgs def reply(self, ser): if ser['_id'] not in self.replied_queue: self.replied_queue.append(ser['_id']) text_out = BotApi(ser['msg']) text_out += ' @%s' % ser['name'] self.api.send_message(text_out, room_id) def _Msgs2Bot(self, msgs: list): """ 解析获取到的聊天记录,筛选出@xbot的记录,保留'_id','name','msg'这三个字段 """ if len(msgs) > 0: df_msgs = pd.DataFrame(msgs) df_msgs['name'] = df_msgs['u'].map(lambda x:x['username']) df_msgs_select = df_msgs[['_id', 'name', 'msg']] self.msgs_at_bot = df_msgs_select[df_msgs_select['msg'].str.startswith('@xbot') | df_msgs_select['msg'].str.endswith('@xbot')] self.msgs_at_bot.apply(self.reply, axis=1) def Msgs2Bot(self): self._Msgs2Bot(self.NewMsgs()) def Run(self): while True: self.Msgs2Bot() time.sleep(0.5)