A Python robot that edits Wikipedia and interacts with people over IRC https://en.wikipedia.org/wiki/User:EarwigBot
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.

358 lines
13 KiB

  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright (C) 2009-2012 by Ben Kurtovic <ben.kurtovic@verizon.net>
  4. #
  5. # Permission is hereby granted, free of charge, to any person obtaining a copy
  6. # of this software and associated documentation files (the "Software"), to deal
  7. # in the Software without restriction, including without limitation the rights
  8. # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  9. # copies of the Software, and to permit persons to whom the Software is
  10. # furnished to do so, subject to the following conditions:
  11. #
  12. # The above copyright notice and this permission notice shall be included in
  13. # all copies or substantial portions of the Software.
  14. #
  15. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16. # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17. # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  18. # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19. # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  20. # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  21. # SOFTWARE.
  22. from getpass import getpass
  23. import logging
  24. import logging.handlers
  25. from os import mkdir, path
  26. import yaml
  27. from earwigbot import blowfish
  28. __all__ = ["BotConfig"]
  29. class BotConfig(object):
  30. """
  31. EarwigBot's YAML Config File Manager
  32. This handles all tasks involving reading and writing to our config file,
  33. including encrypting and decrypting passwords and making a new config file
  34. from scratch at the inital bot run.
  35. BotConfig has a few properties and functions, including the following:
  36. * config.root_dir - bot's working directory; contains config.yml, logs/
  37. * config.path - path to the bot's config file
  38. * config.components - enabled components
  39. * config.wiki - information about wiki-editing
  40. * config.tasks - information for bot tasks
  41. * config.irc - information about IRC
  42. * config.metadata - miscellaneous information
  43. * config.schedule() - tasks scheduled to run at a given time
  44. BotConfig also has some functions used in config loading:
  45. * config.load() - loads and parses our config file, returning True if
  46. passwords are stored encrypted or False otherwise;
  47. can also be used to easily reload config
  48. * config.decrypt() - given a key, decrypts passwords inside our config
  49. variables, and remembers to decrypt the password if
  50. config is reloaded; won't do anything if passwords
  51. aren't encrypted
  52. """
  53. def __init__(self, root_dir, level):
  54. self._root_dir = root_dir
  55. self._logging_level = level
  56. self._config_path = path.join(self._root_dir, "config.yml")
  57. self._log_dir = path.join(self._root_dir, "logs")
  58. self._decryption_key = None
  59. self._data = None
  60. self._components = _ConfigNode()
  61. self._wiki = _ConfigNode()
  62. self._tasks = _ConfigNode()
  63. self._irc = _ConfigNode()
  64. self._metadata = _ConfigNode()
  65. self._nodes = [self._components, self._wiki, self._tasks, self._irc,
  66. self._metadata]
  67. self._decryptable_nodes = [ # Default nodes to decrypt
  68. (self._wiki, ("password")),
  69. (self._wiki, ("search", "credentials", "key")),
  70. (self._wiki, ("search", "credentials", "secret")),
  71. (self._irc, ("frontend", "nickservPassword")),
  72. (self._irc, ("watcher", "nickservPassword")),
  73. ]
  74. def _load(self):
  75. """Load data from our JSON config file (config.yml) into self._data."""
  76. filename = self._config_path
  77. with open(filename, 'r') as fp:
  78. try:
  79. self._data = yaml.load(fp)
  80. except yaml.YAMLError as error:
  81. print "Error parsing config file {0}:".format(filename)
  82. raise
  83. def _setup_logging(self):
  84. """Configures the logging module so it works the way we want it to."""
  85. log_dir = self._log_dir
  86. logger = logging.getLogger("earwigbot")
  87. logger.setLevel(logging.DEBUG)
  88. if self.metadata.get("enableLogging"):
  89. hand = logging.handlers.TimedRotatingFileHandler
  90. formatter = _BotFormatter()
  91. color_formatter = _BotFormatter(color=True)
  92. logfile = lambda f: path.join(log_dir, f)
  93. if not path.isdir(log_dir):
  94. if not path.exists(log_dir):
  95. mkdir(log_dir, 0700)
  96. else:
  97. msg = "log_dir ({0}) exists but is not a directory!"
  98. print msg.format(log_dir)
  99. return
  100. main_handler = hand(logfile("bot.log"), "midnight", 1, 7)
  101. error_handler = hand(logfile("error.log"), "W6", 1, 4)
  102. debug_handler = hand(logfile("debug.log"), "H", 1, 6)
  103. main_handler.setLevel(logging.INFO)
  104. error_handler.setLevel(logging.WARNING)
  105. debug_handler.setLevel(logging.DEBUG)
  106. for h in (main_handler, error_handler, debug_handler):
  107. h.setFormatter(formatter)
  108. logger.addHandler(h)
  109. self._stream_handler = stream = logging.StreamHandler()
  110. stream.setLevel(self._logging_level)
  111. stream.setFormatter(color_formatter)
  112. logger.addHandler(stream)
  113. def _decrypt(self, node, nodes):
  114. """Try to decrypt the contents of a config node. Use self.decrypt()."""
  115. try:
  116. node._decrypt(self._decryption_key, nodes[:-1], nodes[-1])
  117. except blowfish.BlowfishError as error:
  118. print "Error decrypting passwords:"
  119. raise
  120. def _make_new(self):
  121. """Make a new config file based on the user's input."""
  122. #m = "Would you like to encrypt passwords stored in config.yml? [y/n] "
  123. #encrypt = raw_input(m)
  124. #if encrypt.lower().startswith("y"):
  125. # is_encrypted = True
  126. #else:
  127. # is_encrypted = False
  128. raise NotImplementedError()
  129. # yaml.dumps()
  130. @property
  131. def root_dir(self):
  132. return self._root_dir
  133. @property
  134. def logging_level(self):
  135. return self._logging_level
  136. @logging_level.setter
  137. def logging_level(self, level):
  138. self._logging_level = level
  139. self._stream_handler.setLevel(level)
  140. @property
  141. def path(self):
  142. return self._config_path
  143. @property
  144. def log_dir(self):
  145. return self._log_dir
  146. @property
  147. def data(self):
  148. """The entire config file."""
  149. return self._data
  150. @property
  151. def components(self):
  152. """A dict of enabled components."""
  153. return self._components
  154. @property
  155. def wiki(self):
  156. """A dict of information about wiki-editing."""
  157. return self._wiki
  158. @property
  159. def tasks(self):
  160. """A dict of information for bot tasks."""
  161. return self._tasks
  162. @property
  163. def irc(self):
  164. """A dict of information about IRC."""
  165. return self._irc
  166. @property
  167. def metadata(self):
  168. """A dict of miscellaneous information."""
  169. return self._metadata
  170. def is_loaded(self):
  171. """Return True if our config file has been loaded, otherwise False."""
  172. return self._data is not None
  173. def is_encrypted(self):
  174. """Return True if passwords are encrypted, otherwise False."""
  175. return self.metadata.get("encryptPasswords", False)
  176. def load(self):
  177. """Load, or reload, our config file.
  178. First, check if we have a valid config file, and if not, notify the
  179. user. If there is no config file at all, offer to make one, otherwise
  180. exit.
  181. Store data from our config file in five _ConfigNodes (components,
  182. wiki, tasks, irc, metadata) for easy access (as well as the internal
  183. _data variable).
  184. If config is being reloaded, encrypted items will be automatically
  185. decrypted if they were decrypted beforehand.
  186. """
  187. if not path.exists(self._config_path):
  188. print "Config file not found:", self._config_path
  189. choice = raw_input("Would you like to create a config file now? [y/n] ")
  190. if choice.lower().startswith("y"):
  191. self._make_new()
  192. else:
  193. exit(1) # TODO: raise an exception instead
  194. self._load()
  195. data = self._data
  196. self.components._load(data.get("components", {}))
  197. self.wiki._load(data.get("wiki", {}))
  198. self.tasks._load(data.get("tasks", {}))
  199. self.irc._load(data.get("irc", {}))
  200. self.metadata._load(data.get("metadata", {}))
  201. self._setup_logging()
  202. if self.is_encrypted():
  203. if not self._decryption_key:
  204. key = getpass("Enter key to decrypt bot passwords: ")
  205. self._decryption_key = key
  206. for node, nodes in self._decryptable_nodes:
  207. self._decrypt(node, nodes)
  208. def decrypt(self, node, *nodes):
  209. """Use self._decryption_key to decrypt an object in our config tree.
  210. If this is called when passwords are not encrypted (check with
  211. config.is_encrypted()), nothing will happen. We'll also keep track of
  212. this node if config.load() is called again (i.e. to reload) and
  213. automatically decrypt it.
  214. Example usage:
  215. config.decrypt(config.irc, "frontend", "nickservPassword")
  216. -> decrypts config.irc["frontend"]["nickservPassword"]
  217. """
  218. self._decryptable_nodes.append((node, nodes))
  219. if self.is_encrypted():
  220. self._decrypt(node, nodes)
  221. def schedule(self, minute, hour, month_day, month, week_day):
  222. """Return a list of tasks scheduled to run at the specified time.
  223. The schedule data comes from our config file's 'schedule' field, which
  224. is stored as self._data["schedule"]. Call this function as
  225. config.schedule(args).
  226. """
  227. # Tasks to run this turn, each as a list of either [task_name, kwargs],
  228. # or just the task_name:
  229. tasks = []
  230. now = {"minute": minute, "hour": hour, "month_day": month_day,
  231. "month": month, "week_day": week_day}
  232. data = self._data.get("schedule", [])
  233. for event in data:
  234. do = True
  235. for key, value in now.items():
  236. try:
  237. requirement = event[key]
  238. except KeyError:
  239. continue
  240. if requirement != value:
  241. do = False
  242. break
  243. if do:
  244. try:
  245. tasks.extend(event["tasks"])
  246. except KeyError:
  247. pass
  248. return tasks
  249. class _ConfigNode(object):
  250. def __iter__(self):
  251. for key in self.__dict__.iterkeys():
  252. yield key
  253. def __getitem__(self, item):
  254. return self.__dict__.__getitem__(item)
  255. def _dump(self):
  256. data = self.__dict__.copy()
  257. for key, val in data.iteritems():
  258. if isinstance(val, _ConfigNode):
  259. data[key] = val._dump()
  260. return data
  261. def _load(self, data):
  262. self.__dict__ = data.copy()
  263. def _decrypt(self, key, intermediates, item):
  264. base = self.__dict__
  265. try:
  266. for inter in intermediates:
  267. base = base[inter]
  268. except KeyError:
  269. return
  270. if item in base:
  271. base[item] = blowfish.decrypt(key, base[item])
  272. def get(self, *args, **kwargs):
  273. return self.__dict__.get(*args, **kwargs)
  274. class _BotFormatter(logging.Formatter):
  275. def __init__(self, color=False):
  276. self._format = super(_BotFormatter, self).format
  277. if color:
  278. fmt = "[%(asctime)s %(lvl)s] %(name)s: %(message)s"
  279. self.format = lambda record: self._format(self.format_color(record))
  280. else:
  281. fmt = "[%(asctime)s %(levelname)-8s] %(name)s: %(message)s"
  282. self.format = self._format
  283. datefmt = "%Y-%m-%d %H:%M:%S"
  284. super(_BotFormatter, self).__init__(fmt=fmt, datefmt=datefmt)
  285. def format_color(self, record):
  286. l = record.levelname.ljust(8)
  287. if record.levelno == logging.DEBUG:
  288. record.lvl = l.join(("\x1b[34m", "\x1b[0m")) # Blue
  289. if record.levelno == logging.INFO:
  290. record.lvl = l.join(("\x1b[32m", "\x1b[0m")) # Green
  291. if record.levelno == logging.WARNING:
  292. record.lvl = l.join(("\x1b[33m", "\x1b[0m")) # Yellow
  293. if record.levelno == logging.ERROR:
  294. record.lvl = l.join(("\x1b[31m", "\x1b[0m")) # Red
  295. if record.levelno == logging.CRITICAL:
  296. record.lvl = l.join(("\x1b[1m\x1b[31m", "\x1b[0m")) # Bold red
  297. return record