-
Notifications
You must be signed in to change notification settings - Fork 0
/
universal_bot.py
113 lines (94 loc) · 4.33 KB
/
universal_bot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import os
import time
import re
import multiprocessing
from concurrent.futures import ThreadPoolExecutor
from slackclient import SlackClient
from utilities.locks import Lock
from utilities.pending_actions_controller import PendingActionsController
RTM_READ_DELAY = 1
MENTION_REGEX = "^<@(|[WU].+?)>(.*)"
slack_client = SlackClient(os.environ.get('SLACK_BOT_TOKEN'))
universal_bot_id = None
commands = []
executor = ThreadPoolExecutor(max_workers=multiprocessing.cpu_count())
pa_controller = PendingActionsController()
def add_commands(command_classes):
command_names = set()
for item in command_classes:
if item.command_name not in command_names:
command_names.add(item.command_name)
else:
raise RuntimeError("Duplicate of \"{}\" command name found".format(item.command_name))
commands.append((item.command_keys, item))
print("Commands: {} were added".format(", ".join(command_names).join(("(", ")"))))
def process_events(slack_events, multithreaded):
for event in slack_events:
message_for_bot = None
if event["type"] == "message" and "subtype" not in event:
res = slack_client.api_call("conversations.info", channel=event["channel"])
if res["ok"]:
channel_info = res["channel"]
if channel_info["is_im"]:
user_id, message = parse_direct_mention(event["text"])
if user_id == universal_bot_id:
message_for_bot = message
else:
message_for_bot = event["text"]
else:
user_id, message = parse_direct_mention(event["text"])
if user_id == universal_bot_id:
message_for_bot = message
if message_for_bot is None:
continue
res = [x for x in message_for_bot.split(" ") if x]
if len(res) != 0:
if multithreaded:
executor.submit(handle_command, res, event["user"], event["channel"])
else:
handle_command(res, event["user"], event["channel"])
def process_pending():
if pa_controller.is_empty():
return
for keys, command_class in commands:
res = pa_controller.get_actions(command_class.command_name)
if len(res) != 0:
for pending_action in res:
if pending_action.check():
slack_client.api_call(**pending_action.get_action())
pa_controller.remove_action(pending_action)
print("Pending action of \"{}\" was executed".format(command_class.command_name))
def parse_direct_mention(message_text):
matches = re.search(MENTION_REGEX, message_text)
return (matches.group(1), matches.group(2).strip()) if matches else (None, None)
def handle_command(arguments, user_id, channel_id):
for keys, command_class in commands:
if len(arguments) >= len(keys) and arguments[:len(keys)] == keys:
command_instance = command_class(arguments[len(keys):], user_id, channel_id)
command_instance.execute()
actions = command_instance.get_actions()
with Lock():
for action in actions:
slack_client.api_call(**action)
pending_actions = command_instance.get_pending_actions()
print("Command \"{}\" processed by {} in {}".format(command_instance.command_name, user_id, channel_id))
if len(pending_actions) != 0:
pa_controller.add_actions(pending_actions)
if len(pending_actions) == 1:
print("\tand 1 pending action was added")
else:
print("\tand {} pending actions were added".format(len(pending_actions)))
return
def connect():
global universal_bot_id
if not slack_client.rtm_connect(with_team_state=False):
raise RuntimeError("Connection can not be established")
print("Connection established")
universal_bot_id = slack_client.api_call("auth.test")["user_id"]
def loop(multithreaded=False):
Lock.multithreaded = multithreaded
print("Loop started {} multi-threading".format("with" if multithreaded else "without"))
while True:
process_events(slack_client.rtm_read(), multithreaded)
process_pending()
time.sleep(RTM_READ_DELAY)