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.

functions.py 7.4 KiB

13 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. # -*- coding: utf-8 -*-
  2. """
  3. EarwigBot's Wiki Toolset: Misc Functions
  4. This module, a component of the wiki package, contains miscellaneous functions
  5. that are not methods of any class, like get_site().
  6. There's no need to import this module explicitly. All functions here are
  7. automatically available from wiki.
  8. """
  9. from cookielib import LWPCookieJar, LoadError
  10. import errno
  11. from getpass import getpass
  12. from os import chmod, path
  13. import platform
  14. import stat
  15. import config
  16. from wiki.exceptions import SiteNotFoundError
  17. from wiki.site import Site
  18. __all__ = ["get_site"]
  19. _cookiejar = None
  20. def _load_config():
  21. """Called by a config-requiring function, such as get_site(), when config
  22. has not been loaded. This will usually happen only if we're running code
  23. directly from Python's interpreter and not the bot itself, because
  24. earwigbot.py or core/main.py will already call these functions.
  25. """
  26. is_encrypted = config.load()
  27. if is_encrypted: # Passwords in the config file are encrypted
  28. key = getpass("Enter key to unencrypt bot passwords: ")
  29. config.decrypt(key)
  30. def _get_cookiejar():
  31. """Returns a LWPCookieJar object loaded from our .cookies file. The same
  32. one is returned every time.
  33. The .cookies file is located in the project root, same directory as
  34. config.json and earwigbot.py. If it doesn't exist, we will create the file
  35. and set it to be readable and writeable only by us. If it exists but the
  36. information inside is bogus, we will ignore it.
  37. This is normally called by _get_site_object_from_dict() (in turn called by
  38. get_site()), and the cookiejar is passed to our Site's constructor, used
  39. when it makes API queries. This way, we can easily preserve cookies between
  40. sites (e.g., for CentralAuth), making logins easier.
  41. """
  42. global _cookiejar
  43. if _cookiejar is not None:
  44. return _cookiejar
  45. cookie_file = path.join(config.root_dir, ".cookies")
  46. _cookiejar = LWPCookieJar(cookie_file)
  47. try:
  48. _cookiejar.load()
  49. except LoadError:
  50. pass # File contains bad data, so ignore it completely
  51. except IOError as e:
  52. if e.errno == errno.ENOENT: # "No such file or directory"
  53. # Create the file and restrict reading/writing only to the owner,
  54. # so others can't peak at our cookies:
  55. open(cookie_file, "w").close()
  56. chmod(cookie_file, stat.S_IRUSR|stat.S_IWUSR)
  57. else:
  58. raise
  59. return _cookiejar
  60. def _get_site_object_from_dict(name, d):
  61. """Return a Site object based on the contents of a dict, probably acquired
  62. through our config file, and a separate name.
  63. """
  64. project = d.get("project")
  65. lang = d.get("lang")
  66. base_url = d.get("baseURL")
  67. article_path = d.get("articlePath")
  68. script_path = d.get("scriptPath")
  69. sql = d.get("sql", {})
  70. namespaces = d.get("namespaces", {})
  71. login = (config.wiki.get("username"), config.wiki.get("password"))
  72. cookiejar = _get_cookiejar()
  73. user_agent = config.wiki.get("userAgent")
  74. assert_edit = config.wiki.get("assert")
  75. maxlag = config.wiki.get("maxlag")
  76. if user_agent:
  77. user_agent = user_agent.replace("$1", platform.python_version())
  78. for key, value in namespaces.items(): # Convert string keys to integers
  79. del namespaces[key]
  80. try:
  81. namespaces[int(key)] = value
  82. except ValueError: # Data is broken, ignore it
  83. namespaces = None
  84. break
  85. return Site(name=name, project=project, lang=lang, base_url=base_url,
  86. article_path=article_path, script_path=script_path, sql=sql,
  87. namespaces=namespaces, login=login, cookiejar=cookiejar,
  88. user_agent=user_agent, assert_edit=assert_edit, maxlag=maxlag)
  89. def get_site(name=None, project=None, lang=None):
  90. """Returns a Site instance based on information from our config file.
  91. With no arguments, returns the default site as specified by our config
  92. file. This is default = config.wiki["defaultSite"];
  93. config.wiki["sites"][default].
  94. With `name` specified, returns the site specified by
  95. config.wiki["sites"][name].
  96. With `project` and `lang` specified, returns the site specified by the
  97. member of config.wiki["sites"], `s`, for which s["project"] == project and
  98. s["lang"] == lang.
  99. We will attempt to login to the site automatically
  100. using config.wiki["username"] and config.wiki["password"] if both are
  101. defined.
  102. Specifying a project without a lang or a lang without a project will raise
  103. TypeError. If all three args are specified, `name` will be first tried,
  104. then `project` and `lang`. If, with any number of args, a site cannot be
  105. found in the config, SiteNotFoundError is raised.
  106. """
  107. # Check if config has been loaded, and load it if it hasn't:
  108. if not config.is_loaded():
  109. _load_config()
  110. # Someone specified a project without a lang (or a lang without a project)!
  111. if (project is None and lang is not None) or (project is not None and
  112. lang is None):
  113. e = "Keyword arguments 'lang' and 'project' must be specified together."
  114. raise TypeError(e)
  115. # No args given, so return our default site (project is None implies lang
  116. # is None, so we don't need to add that in):
  117. if name is None and project is None:
  118. try:
  119. default = config.wiki["defaultSite"]
  120. except KeyError:
  121. e = "Default site is not specified in config."
  122. raise SiteNotFoundError(e)
  123. try:
  124. site = config.wiki["sites"][default]
  125. except KeyError:
  126. e = "Default site specified by config is not in the config's sites list."
  127. raise SiteNotFoundError(e)
  128. return _get_site_object_from_dict(default, site)
  129. # Name arg given, but don't look at others unless `name` isn't found:
  130. if name is not None:
  131. try:
  132. site = config.wiki["sites"][name]
  133. except KeyError:
  134. if project is None: # Implies lang is None, so only name was given
  135. e = "Site '{0}' not found in config.".format(name)
  136. raise SiteNotFoundError(e)
  137. for sitename, site in config.wiki["sites"].items():
  138. if site["project"] == project and site["lang"] == lang:
  139. return _get_site_object_from_dict(sitename, site)
  140. e = "Neither site '{0}' nor site '{1}:{2}' found in config."
  141. e.format(name, project, lang)
  142. raise SiteNotFoundError(e)
  143. else:
  144. return _get_site_object_from_dict(name, site)
  145. # If we end up here, then project and lang are both not None:
  146. for sitename, site in config.wiki["sites"].items():
  147. if site["project"] == project and site["lang"] == lang:
  148. return _get_site_object_from_dict(sitename, site)
  149. e = "Site '{0}:{1}' not found in config.".format(project, lang)
  150. raise SiteNotFoundError(e)
  151. def add_site():
  152. """STUB: config editing is required first.
  153. Returns True if the site was added successfully or False if the site was
  154. already in our config. Raises ConfigError if saving the updated file failed
  155. for some reason."""
  156. pass
  157. def del_site(name):
  158. """STUB: config editing is required first.
  159. Returns True if the site was removed successfully or False if the site was
  160. not in our config originally. Raises ConfigError if saving the updated file
  161. failed for some reason."""
  162. pass