bot.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417
  1. # -*- coding: utf-8 -*-
  2. """
  3. Simple IRC bot library.
  4. This module contains a single-server IRC bot class that can be used to
  5. write simpler bots.
  6. """
  7. from __future__ import absolute_import
  8. import sys
  9. import collections
  10. import itertools
  11. import irc.client
  12. import irc.modes
  13. from .dict import IRCDict
  14. class ServerSpec(object):
  15. """
  16. An IRC server specification.
  17. >>> spec = ServerSpec('localhost')
  18. >>> spec.host
  19. 'localhost'
  20. >>> spec.port
  21. 6667
  22. >>> spec.password
  23. >>> spec = ServerSpec('127.0.0.1', 6697, 'fooP455')
  24. >>> spec.password
  25. 'fooP455'
  26. """
  27. def __init__(self, host, port=6667, password=None):
  28. self.host = host
  29. self.port = port
  30. self.password = password
  31. class SingleServerIRCBot(irc.client.SimpleIRCClient):
  32. """A single-server IRC bot class.
  33. The bot tries to reconnect if it is disconnected.
  34. The bot keeps track of the channels it has joined, the other
  35. clients that are present in the channels and which of those that
  36. have operator or voice modes. The "database" is kept in the
  37. self.channels attribute, which is an IRCDict of Channels.
  38. """
  39. def __init__(self, server_list, nickname, realname,
  40. reconnection_interval=60, **connect_params):
  41. """Constructor for SingleServerIRCBot objects.
  42. Arguments:
  43. server_list -- A list of ServerSpec objects or tuples of
  44. parameters suitable for constructing ServerSpec
  45. objects. Defines the list of servers the bot will
  46. use (in order).
  47. nickname -- The bot's nickname.
  48. realname -- The bot's realname.
  49. reconnection_interval -- How long the bot should wait
  50. before trying to reconnect.
  51. dcc_connections -- A list of initiated/accepted DCC
  52. connections.
  53. **connect_params -- parameters to pass through to the connect
  54. method.
  55. """
  56. super(SingleServerIRCBot, self).__init__()
  57. self.__connect_params = connect_params
  58. self.channels = IRCDict()
  59. self.server_list = [
  60. ServerSpec(*server)
  61. if isinstance(server, (tuple, list))
  62. else server
  63. for server in server_list
  64. ]
  65. assert all(
  66. isinstance(server, ServerSpec)
  67. for server in self.server_list
  68. )
  69. if not reconnection_interval or reconnection_interval < 0:
  70. reconnection_interval = 2 ** 31
  71. self.reconnection_interval = reconnection_interval
  72. self._nickname = nickname
  73. self._realname = realname
  74. for i in ["disconnect", "join", "kick", "mode",
  75. "namreply", "nick", "part", "quit"]:
  76. self.connection.add_global_handler(i, getattr(self, "_on_" + i),
  77. -20)
  78. def _connected_checker(self):
  79. """[Internal]"""
  80. if not self.connection.is_connected():
  81. self.connection.execute_delayed(self.reconnection_interval,
  82. self._connected_checker)
  83. self.jump_server()
  84. def _connect(self):
  85. """
  86. Establish a connection to the server at the front of the server_list.
  87. """
  88. server = self.server_list[0]
  89. try:
  90. self.connect(server.host, server.port, self._nickname,
  91. server.password, ircname=self._realname,
  92. **self.__connect_params)
  93. except irc.client.ServerConnectionError:
  94. pass
  95. def _on_disconnect(self, c, e):
  96. self.channels = IRCDict()
  97. self.connection.execute_delayed(self.reconnection_interval,
  98. self._connected_checker)
  99. def _on_join(self, c, e):
  100. ch = e.target
  101. nick = e.source.nick
  102. if nick == c.get_nickname():
  103. self.channels[ch] = Channel()
  104. self.channels[ch].add_user(nick)
  105. def _on_kick(self, c, e):
  106. nick = e.arguments[0]
  107. channel = e.target
  108. if nick == c.get_nickname():
  109. del self.channels[channel]
  110. else:
  111. self.channels[channel].remove_user(nick)
  112. def _on_mode(self, c, e):
  113. t = e.target
  114. if not irc.client.is_channel(t):
  115. # mode on self; disregard
  116. return
  117. ch = self.channels[t]
  118. modes = irc.modes.parse_channel_modes(" ".join(e.arguments))
  119. for sign, mode, argument in modes:
  120. f = {"+": ch.set_mode, "-": ch.clear_mode}[sign]
  121. f(mode, argument)
  122. def _on_namreply(self, c, e):
  123. """
  124. e.arguments[0] == "@" for secret channels,
  125. "*" for private channels,
  126. "=" for others (public channels)
  127. e.arguments[1] == channel
  128. e.arguments[2] == nick list
  129. """
  130. ch_type, channel, nick_list = e.arguments
  131. if channel == '*':
  132. # User is not in any visible channel
  133. # http://tools.ietf.org/html/rfc2812#section-3.2.5
  134. return
  135. for nick in nick_list.split():
  136. nick_modes = []
  137. if nick[0] in self.connection.features.prefix:
  138. nick_modes.append(self.connection.features.prefix[nick[0]])
  139. nick = nick[1:]
  140. for mode in nick_modes:
  141. self.channels[channel].set_mode(mode, nick)
  142. self.channels[channel].add_user(nick)
  143. def _on_nick(self, c, e):
  144. before = e.source.nick
  145. after = e.target
  146. for ch in self.channels.values():
  147. if ch.has_user(before):
  148. ch.change_nick(before, after)
  149. def _on_part(self, c, e):
  150. nick = e.source.nick
  151. channel = e.target
  152. if nick == c.get_nickname():
  153. del self.channels[channel]
  154. else:
  155. self.channels[channel].remove_user(nick)
  156. def _on_quit(self, c, e):
  157. nick = e.source.nick
  158. for ch in self.channels.values():
  159. if ch.has_user(nick):
  160. ch.remove_user(nick)
  161. def die(self, msg="Bye, cruel world!"):
  162. """Let the bot die.
  163. Arguments:
  164. msg -- Quit message.
  165. """
  166. self.connection.disconnect(msg)
  167. sys.exit(0)
  168. def disconnect(self, msg="I'll be back!"):
  169. """Disconnect the bot.
  170. The bot will try to reconnect after a while.
  171. Arguments:
  172. msg -- Quit message.
  173. """
  174. self.connection.disconnect(msg)
  175. def get_version(self):
  176. """Returns the bot version.
  177. Used when answering a CTCP VERSION request.
  178. """
  179. return "Python irc.bot ({version})".format(
  180. version=irc.client.VERSION_STRING)
  181. def jump_server(self, msg="Changing servers"):
  182. """Connect to a new server, possibly disconnecting from the current.
  183. The bot will skip to next server in the server_list each time
  184. jump_server is called.
  185. """
  186. if self.connection.is_connected():
  187. self.connection.disconnect(msg)
  188. self.server_list.append(self.server_list.pop(0))
  189. self._connect()
  190. def on_ctcp(self, c, e):
  191. """Default handler for ctcp events.
  192. Replies to VERSION and PING requests and relays DCC requests
  193. to the on_dccchat method.
  194. """
  195. nick = e.source.nick
  196. if e.arguments[0] == "VERSION":
  197. c.ctcp_reply(nick, "VERSION " + self.get_version())
  198. elif e.arguments[0] == "PING":
  199. if len(e.arguments) > 1:
  200. c.ctcp_reply(nick, "PING " + e.arguments[1])
  201. elif e.arguments[0] == "DCC" and e.arguments[1].split(" ", 1)[0] == "CHAT":
  202. self.on_dccchat(c, e)
  203. def on_dccchat(self, c, e):
  204. pass
  205. def start(self):
  206. """Start the bot."""
  207. self._connect()
  208. super(SingleServerIRCBot, self).start()
  209. class Channel(object):
  210. """
  211. A class for keeping information about an IRC channel.
  212. """
  213. user_modes = 'ovqha'
  214. """
  215. Modes which are applicable to individual users, and which
  216. should be tracked in the mode_users dictionary.
  217. """
  218. def __init__(self):
  219. self._users = IRCDict()
  220. self.mode_users = collections.defaultdict(IRCDict)
  221. self.modes = {}
  222. def users(self):
  223. """Returns an unsorted list of the channel's users."""
  224. return self._users.keys()
  225. def opers(self):
  226. """Returns an unsorted list of the channel's operators."""
  227. return self.mode_users['o'].keys()
  228. def voiced(self):
  229. """Returns an unsorted list of the persons that have voice
  230. mode set in the channel."""
  231. return self.mode_users['v'].keys()
  232. def owners(self):
  233. """Returns an unsorted list of the channel's owners."""
  234. return self.mode_users['q'].keys()
  235. def halfops(self):
  236. """Returns an unsorted list of the channel's half-operators."""
  237. return self.mode_users['h'].keys()
  238. def admins(self):
  239. """Returns an unsorted list of the channel's admins."""
  240. return self.mode_users['a'].keys()
  241. def has_user(self, nick):
  242. """Check whether the channel has a user."""
  243. return nick in self._users
  244. def is_oper(self, nick):
  245. """Check whether a user has operator status in the channel."""
  246. return nick in self.mode_users['o']
  247. def is_voiced(self, nick):
  248. """Check whether a user has voice mode set in the channel."""
  249. return nick in self.mode_users['v']
  250. def is_owner(self, nick):
  251. """Check whether a user has owner status in the channel."""
  252. return nick in self.mode_users['q']
  253. def is_halfop(self, nick):
  254. """Check whether a user has half-operator status in the channel."""
  255. return nick in self.mode_users['h']
  256. def is_admin(self, nick):
  257. """Check whether a user has admin status in the channel."""
  258. return nick in self.mode_users['a']
  259. def add_user(self, nick):
  260. self._users[nick] = 1
  261. @property
  262. def user_dicts(self):
  263. yield self._users
  264. for d in self.mode_users.values():
  265. yield d
  266. def remove_user(self, nick):
  267. for d in self.user_dicts:
  268. d.pop(nick, None)
  269. def change_nick(self, before, after):
  270. self._users[after] = self._users.pop(before)
  271. for mode_lookup in self.mode_users.values():
  272. if before in mode_lookup:
  273. mode_lookup[after] = mode_lookup.pop(before)
  274. def set_userdetails(self, nick, details):
  275. if nick in self._users:
  276. self._users[nick] = details
  277. def set_mode(self, mode, value=None):
  278. """Set mode on the channel.
  279. Arguments:
  280. mode -- The mode (a single-character string).
  281. value -- Value
  282. """
  283. if mode in self.user_modes:
  284. self.mode_users[mode][value] = 1
  285. else:
  286. self.modes[mode] = value
  287. def clear_mode(self, mode, value=None):
  288. """Clear mode on the channel.
  289. Arguments:
  290. mode -- The mode (a single-character string).
  291. value -- Value
  292. """
  293. try:
  294. if mode in self.user_modes:
  295. del self.mode_users[mode][value]
  296. else:
  297. del self.modes[mode]
  298. except KeyError:
  299. pass
  300. def has_mode(self, mode):
  301. return mode in self.modes
  302. def is_moderated(self):
  303. return self.has_mode("m")
  304. def is_secret(self):
  305. return self.has_mode("s")
  306. def is_protected(self):
  307. return self.has_mode("p")
  308. def has_topic_lock(self):
  309. return self.has_mode("t")
  310. def is_invite_only(self):
  311. return self.has_mode("i")
  312. def has_allow_external_messages(self):
  313. return self.has_mode("n")
  314. def has_limit(self):
  315. return self.has_mode("l")
  316. def limit(self):
  317. if self.has_limit():
  318. return self.modes["l"]
  319. else:
  320. return None
  321. def has_key(self):
  322. return self.has_mode("k")