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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
|
#!/usr/bin/env python3
# -*- mode: python -*-
import asyncio
import datetime
import functools
import logging.config
import os
from warmachine.config import Config
from warmachine.connections.irc import AioIRC
from warmachine.connections.slack import SlackWS
log_config = {
'version': 1,
'disable_existing_loggers': False,
'formatters': {
'standard': {
'format': '%(asctime)s [%(levelname)s] %(name)s: %(message)s',
}
},
'handlers': {
'console': {
'class': 'logging.StreamHandler',
'level': 'DEBUG',
'formatter': 'standard',
},
},
'loggers': {
'websockets': {
'level': 'INFO',
'handlers': ['console', ]
},
'': {
'level': 'INFO',
'handlers': ['console', ]
}
}
}
class Bot(object):
def __init__(self, settings):
self.log = logging.getLogger(self.__class__.__name__)
self._loop = asyncio.get_event_loop()
self.settings = settings
self.connections = {}
self.tasks = []
self.loaded_plugins = []
self.load_plugin('warmachine.addons.giphy.GiphySearch')
self.load_plugin('warmachine.addons.standup.StandUpPlugin')
# TODO: Ensure the config directory has been created
def start(self):
for connection in self.connections:
t = asyncio.ensure_future(connection.connect())
t.add_done_callback(functools.partial(self.on_connect, connection))
self._loop.run_forever()
def add_connection(self, connection):
self.connections[connection] = {}
def on_connect(self, connection, task):
for p in self.loaded_plugins:
if hasattr(p, 'on_connect'):
p.on_connect(connection)
asyncio.ensure_future(self.process_message(connection))
async def process_message(self, connection):
"""
Constantly read new messages from the connection in a non-blocking way
"""
while True:
message = await connection.read()
if not message:
continue
for p in self.loaded_plugins:
self.log.debug('Calling {}'.format(p.__class__.__name__))
try:
await p.recv_msg(connection, message)
except Exception as e:
self.log.exception(e)
continue
self.log.debug('MSG {}: {}'.format(
connection.__class__.__name__, message))
def load_plugin(self, class_path):
"""
Loads plugins
"""
from importlib import import_module
mod_path, cls_name = class_path.rsplit('.', 1)
mod = import_module(mod_path)
if hasattr(mod, cls_name):
cls = getattr(mod, cls_name)()
self.loaded_plugins.append(cls)
def reload_plugin(self, path):
"""
Reload a plugin
"""
def unload_plugin(self, path):
"""
Unload a plugin
"""
if __name__ == "__main__":
import argparse
import sys
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', help='define warmachine config file',
type=str)
parser.add_argument('--debug', help='enable extra logging output',
action='store_true', default=False)
args = parser.parse_args()
if args.config:
settings = Config(args.config)
else:
sys.stderr.write('Please specify a config file\n')
sys.exit(1)
if args.debug:
log_config['loggers']['']['level'] = 'DEBUG'
logging.config.dictConfig(log_config)
bot = Bot(settings)
for s in settings.sections():
options = settings.options_as_dict(s)
if options.get('enable', False) == 'true':
if s.startswith('slack'):
bot.add_connection(SlackWS(options))
# elif s.startswith('irc'):
# bot.add_connection(AioIRC(options))
bot.start()
|
