A Python robot that edits Wikipedia and interacts with people over IRC https://en.wikipedia.org/wiki/User:EarwigBot
Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

vor 13 Jahren
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610
  1. # -*- coding: utf-8 -*-
  2. from cookielib import CookieJar
  3. from gzip import GzipFile
  4. from json import loads
  5. from os.path import expanduser
  6. from re import escape as re_escape, match as re_match
  7. from StringIO import StringIO
  8. from time import sleep
  9. from urllib import unquote_plus, urlencode
  10. from urllib2 import build_opener, HTTPCookieProcessor, URLError
  11. from urlparse import urlparse
  12. try:
  13. import oursql
  14. except ImportError:
  15. oursql = None
  16. from wiki.category import Category
  17. from wiki.constants import *
  18. from wiki.exceptions import *
  19. from wiki.page import Page
  20. from wiki.user import User
  21. class Site(object):
  22. """
  23. EarwigBot's Wiki Toolset: Site Class
  24. Represents a Site, with support for API queries and returning Pages, Users,
  25. and Categories. The constructor takes a bunch of arguments and you probably
  26. won't need to call it directly, rather tools.get_site() for returning Site
  27. instances, tools.add_site() for adding new ones to config, and
  28. tools.del_site() for removing old ones from config, should suffice.
  29. Public methods:
  30. name -- returns our name (or "wikiid"), like "enwiki"
  31. project -- returns our project name, like "wikipedia"
  32. lang -- returns our language code, like "en"
  33. domain -- returns our web domain, like "en.wikipedia.org"
  34. api_query -- does an API query with the given kwargs as params
  35. sql_query -- does an SQL query and yields its results
  36. get_replag -- returns the estimated database replication lag
  37. namespace_id_to_name -- given a namespace ID, returns associated name(s)
  38. namespace_name_to_id -- given a namespace name, returns associated id
  39. get_page -- returns a Page object for the given title
  40. get_category -- returns a Category object for the given title
  41. get_user -- returns a User object for the given username
  42. """
  43. def __init__(self, name=None, project=None, lang=None, base_url=None,
  44. article_path=None, script_path=None, sql=None,
  45. namespaces=None, login=(None, None), cookiejar=None,
  46. user_agent=None, assert_edit=None, maxlag=None):
  47. """Constructor for new Site instances.
  48. This probably isn't necessary to call yourself unless you're building a
  49. Site that's not in your config and you don't want to add it - normally
  50. all you need is tools.get_site(name), which creates the Site for you
  51. based on your config file. We accept a bunch of kwargs, but the only
  52. ones you really "need" are `base_url` and `script_path` - this is
  53. enough to figure out an API url. `login`, a tuple of
  54. (username, password), is highly recommended. `cookiejar` will be used
  55. to store cookies, and we'll use a normal CookieJar if none is given.
  56. First, we'll store the given arguments as attributes, then set up our
  57. URL opener. We'll load any of the attributes that weren't given from
  58. the API, and then log in if a username/pass was given and we aren't
  59. already logged in.
  60. """
  61. # Attributes referring to site information, filled in by an API query
  62. # if they are missing (and an API url can be determined):
  63. self._name = name
  64. self._project = project
  65. self._lang = lang
  66. self._base_url = base_url
  67. self._article_path = article_path
  68. self._script_path = script_path
  69. self._namespaces = namespaces
  70. # Attributes used for API queries:
  71. self._assert_edit = assert_edit
  72. self._maxlag = maxlag
  73. self._max_retries = 5
  74. # Attributes used for SQL queries:
  75. self._sql_data = sql
  76. self._sql_conn = None
  77. # Set up cookiejar and URL opener for making API queries:
  78. if cookiejar is not None:
  79. self._cookiejar = cookiejar
  80. else:
  81. self._cookiejar = CookieJar()
  82. if user_agent is None:
  83. user_agent = USER_AGENT # Set default UA from wiki.constants
  84. self._opener = build_opener(HTTPCookieProcessor(self._cookiejar))
  85. self._opener.addheaders = [("User-Agent", user_agent),
  86. ("Accept-Encoding", "gzip")]
  87. # Get all of the above attributes that were not specified as arguments:
  88. self._load_attributes()
  89. # If we have a name/pass and the API says we're not logged in, log in:
  90. self._login_info = name, password = login
  91. if name is not None and password is not None:
  92. logged_in_as = self._get_username_from_cookies()
  93. if logged_in_as is None or name != logged_in_as:
  94. self._login(login)
  95. def __repr__(self):
  96. """Returns the canonical string representation of the Site."""
  97. res = ", ".join((
  98. "Site(name={_name!r}", "project={_project!r}", "lang={_lang!r}",
  99. "base_url={_base_url!r}", "article_path={_article_path!r}",
  100. "script_path={_script_path!r}", "assert_edit={_assert_edit!r}",
  101. "maxlag={_maxlag!r}", "sql={_sql!r}", "login={0}",
  102. "user_agent={2!r}", "cookiejar={1})"
  103. ))
  104. name, password = self._login_info
  105. login = "({0}, {1})".format(repr(name), "hidden" if password else None)
  106. cookies = self._cookiejar.__class__.__name__
  107. try:
  108. cookies += "({0!r})".format(self._cookiejar.filename)
  109. except AttributeError:
  110. cookies += "()"
  111. agent = self._opener.addheaders[0][1]
  112. return res.format(login, cookies, agent, **self.__dict__)
  113. def __str__(self):
  114. """Returns a nice string representation of the Site."""
  115. res = "<Site {0} ({1}:{2}) at {3}>"
  116. return res.format(self.name(), self.project(), self.lang(),
  117. self.domain())
  118. def _api_query(self, params, tries=0, wait=5):
  119. """Do an API query with `params` as a dict of parameters.
  120. This will first attempt to construct an API url from self._base_url and
  121. self._script_path. We need both of these, or else we'll raise
  122. SiteAPIError.
  123. We'll encode the given params, adding format=json along the way, as
  124. well as &assert= and &maxlag= based on self._assert_edit and _maxlag.
  125. We make the request through self._opener, which has built-in cookie
  126. support via self._cookiejar, a User-Agent (wiki.constants.USER_AGENT),
  127. and Accept-Encoding set to "gzip".
  128. Assuming everything went well, we'll gunzip the data (if compressed),
  129. load it as a JSON object, and return it.
  130. If our request failed for some reason, we'll raise SiteAPIError with
  131. details. If that reason was due to maxlag, we'll sleep for a bit and
  132. then repeat the query until we exceed self._max_retries.
  133. There's helpful MediaWiki API documentation at
  134. <http://www.mediawiki.org/wiki/API>.
  135. """
  136. if self._base_url is None or self._script_path is None:
  137. e = "Tried to do an API query, but no API URL is known."
  138. raise SiteAPIError(e)
  139. url = ''.join((self._base_url, self._script_path, "/api.php"))
  140. params["format"] = "json" # This is the only format we understand
  141. if self._assert_edit: # If requested, ensure that we're logged in
  142. params["assert"] = self._assert_edit
  143. if self._maxlag: # If requested, don't overload the servers
  144. params["maxlag"] = self._maxlag
  145. data = urlencode(params)
  146. print url, data # debug code
  147. try:
  148. response = self._opener.open(url, data)
  149. except URLError as error:
  150. if hasattr(error, "reason"):
  151. e = "API query failed: {0}.".format(error.reason)
  152. elif hasattr(error, "code"):
  153. e = "API query failed: got an error code of {0}."
  154. e = e.format(error.code)
  155. else:
  156. e = "API query failed."
  157. raise SiteAPIError(e)
  158. result = response.read()
  159. if response.headers.get("Content-Encoding") == "gzip":
  160. stream = StringIO(result)
  161. gzipper = GzipFile(fileobj=stream)
  162. result = gzipper.read()
  163. try:
  164. res = loads(result) # Parse as a JSON object
  165. except ValueError:
  166. e = "API query failed: JSON could not be decoded."
  167. raise SiteAPIError(e)
  168. try:
  169. code = res["error"]["code"]
  170. info = res["error"]["info"]
  171. except (TypeError, KeyError):
  172. return res
  173. if code == "maxlag":
  174. if tries >= self._max_retries:
  175. e = "Maximum number of retries reached ({0})."
  176. raise SiteAPIError(e.format(self._max_retries))
  177. tries += 1
  178. msg = 'Server says: "{0}". Retrying in {1} seconds ({2}/{3}).'
  179. print msg.format(info, wait, tries, self._max_retries)
  180. sleep(wait)
  181. return self._api_query(params, tries=tries, wait=wait*3)
  182. else:
  183. e = 'API query failed: got error "{0}"; server says: "{1}".'
  184. error = SiteAPIError(e.format(code, info))
  185. error.code, error.info = code, info
  186. raise error
  187. def _load_attributes(self, force=False):
  188. """Load data about our Site from the API.
  189. This function is called by __init__() when one of the site attributes
  190. was not given as a keyword argument. We'll do an API query to get the
  191. missing data, but only if there actually *is* missing data.
  192. Additionally, you can call this with `force=True` to forcibly reload
  193. all attributes.
  194. """
  195. # All attributes to be loaded, except _namespaces, which is a special
  196. # case because it requires additional params in the API query:
  197. attrs = [self._name, self._project, self._lang, self._base_url,
  198. self._article_path, self._script_path]
  199. params = {"action": "query", "meta": "siteinfo"}
  200. if not self._namespaces or force:
  201. params["siprop"] = "general|namespaces|namespacealiases"
  202. result = self._api_query(params)
  203. self._load_namespaces(result)
  204. elif all(attrs): # Everything is already specified and we're not told
  205. return # to force a reload, so do nothing
  206. else: # We're only loading attributes other than _namespaces
  207. params["siprop"] = "general"
  208. result = self._api_query(params)
  209. res = result["query"]["general"]
  210. self._name = res["wikiid"]
  211. self._project = res["sitename"].lower()
  212. self._lang = res["lang"]
  213. self._base_url = res["server"]
  214. self._article_path = res["articlepath"]
  215. self._script_path = res["scriptpath"]
  216. def _load_namespaces(self, result):
  217. """Fill self._namespaces with a dict of namespace IDs and names.
  218. Called by _load_attributes() with API data as `result` when
  219. self._namespaces was not given as an kwarg to __init__().
  220. """
  221. self._namespaces = {}
  222. for namespace in result["query"]["namespaces"].values():
  223. ns_id = namespace["id"]
  224. name = namespace["*"]
  225. try:
  226. canonical = namespace["canonical"]
  227. except KeyError:
  228. self._namespaces[ns_id] = [name]
  229. else:
  230. if name != canonical:
  231. self._namespaces[ns_id] = [name, canonical]
  232. else:
  233. self._namespaces[ns_id] = [name]
  234. for namespace in result["query"]["namespacealiases"]:
  235. ns_id = namespace["id"]
  236. alias = namespace["*"]
  237. self._namespaces[ns_id].append(alias)
  238. def _get_cookie(self, name, domain):
  239. """Return the named cookie unless it is expired or doesn't exist."""
  240. for cookie in self._cookiejar:
  241. if cookie.name == name and cookie.domain == domain:
  242. if cookie.is_expired():
  243. break
  244. return cookie
  245. def _get_username_from_cookies(self):
  246. """Try to return our username based solely on cookies.
  247. First, we'll look for a cookie named self._name + "Token", like
  248. "enwikiToken". If it exists and isn't expired, we'll assume it's valid
  249. and try to return the value of the cookie self._name + "UserName" (like
  250. "enwikiUserName"). This should work fine on wikis without single-user
  251. login.
  252. If `enwikiToken` doesn't exist, we'll try to find a cookie named
  253. `centralauth_Token`. If this exists and is not expired, we'll try to
  254. return the value of `centralauth_User`.
  255. If we didn't get any matches, we'll return None. Our goal here isn't to
  256. return the most likely username, or what we *want* our username to be
  257. (for that, we'd do self._login_info[0]), but rather to get our current
  258. username without an unnecessary ?action=query&meta=userinfo API query.
  259. """
  260. domain = self.domain()
  261. name = ''.join((self._name, "Token"))
  262. cookie = self._get_cookie(name, domain)
  263. if cookie is not None:
  264. name = ''.join((self._name, "UserName"))
  265. user_name = self._get_cookie(name, domain)
  266. if user_name is not None:
  267. return user_name.value
  268. name = "centralauth_Token"
  269. for cookie in self._cookiejar:
  270. if cookie.domain_initial_dot is False or cookie.is_expired():
  271. continue
  272. if cookie.name != name:
  273. continue
  274. # Build a regex that will match domains this cookie affects:
  275. search = ''.join(("(.*?)", re_escape(cookie.domain)))
  276. if re_match(search, domain): # Test it against our site
  277. user_name = self._get_cookie("centralauth_User", cookie.domain)
  278. if user_name is not None:
  279. return user_name.value
  280. def _get_username_from_api(self):
  281. """Do a simple API query to get our username and return it.
  282. This is a reliable way to make sure we are actually logged in, because
  283. it doesn't deal with annoying cookie logic, but it results in an API
  284. query that is unnecessary in some cases.
  285. Called by _get_username() (in turn called by get_user() with no
  286. username argument) when cookie lookup fails, probably indicating that
  287. we are logged out.
  288. """
  289. params = {"action": "query", "meta": "userinfo"}
  290. result = self._api_query(params)
  291. return result["query"]["userinfo"]["name"]
  292. def _get_username(self):
  293. """Return the name of the current user, whether logged in or not.
  294. First, we'll try to deduce it solely from cookies, to avoid an
  295. unnecessary API query. For the cookie-detection method, see
  296. _get_username_from_cookies()'s docs.
  297. If our username isn't in cookies, then we're probably not logged in, or
  298. something fishy is going on (like forced logout). In this case, do a
  299. single API query for our username (or IP address) and return that.
  300. """
  301. name = self._get_username_from_cookies()
  302. if name is not None:
  303. return name
  304. return self._get_username_from_api()
  305. def _save_cookiejar(self):
  306. """Try to save our cookiejar after doing a (normal) login or logout.
  307. Calls the standard .save() method with no filename. Don't fret if our
  308. cookiejar doesn't support saving (CookieJar raises AttributeError,
  309. FileCookieJar raises NotImplementedError) or no default filename was
  310. given (LWPCookieJar and MozillaCookieJar raise ValueError).
  311. """
  312. try:
  313. self._cookiejar.save()
  314. except (AttributeError, NotImplementedError, ValueError):
  315. pass
  316. def _login(self, login, token=None, attempt=0):
  317. """Safely login through the API.
  318. Normally, this is called by __init__() if a username and password have
  319. been provided and no valid login cookies were found. The only other
  320. time it needs to be called is when those cookies expire, which is done
  321. automatically by api_query() if a query fails.
  322. Recent versions of MediaWiki's API have fixed a CSRF vulnerability,
  323. requiring login to be done in two separate requests. If the response
  324. from from our initial request is "NeedToken", we'll do another one with
  325. the token. If login is successful, we'll try to save our cookiejar.
  326. Raises LoginError on login errors (duh), like bad passwords and
  327. nonexistent usernames.
  328. `login` is a (username, password) tuple. `token` is the token returned
  329. from our first request, and `attempt` is to prevent getting stuck in a
  330. loop if MediaWiki isn't acting right.
  331. """
  332. name, password = login
  333. params = {"action": "login", "lgname": name, "lgpassword": password}
  334. if token is not None:
  335. params["lgtoken"] = token
  336. result = self._api_query(params)
  337. res = result["login"]["result"]
  338. if res == "Success":
  339. self._save_cookiejar()
  340. elif res == "NeedToken" and attempt == 0:
  341. token = result["login"]["token"]
  342. return self._login(login, token, attempt=1)
  343. else:
  344. if res == "Illegal":
  345. e = "The provided username is illegal."
  346. elif res == "NotExists":
  347. e = "The provided username does not exist."
  348. elif res == "EmptyPass":
  349. e = "No password was given."
  350. elif res == "WrongPass" or res == "WrongPluginPass":
  351. e = "The given password is incorrect."
  352. else:
  353. e = "Couldn't login; server says '{0}'.".format(res)
  354. raise LoginError(e)
  355. def _logout(self):
  356. """Safely logout through the API.
  357. We'll do a simple API request (api.php?action=logout), clear our
  358. cookiejar (which probably contains now-invalidated cookies) and try to
  359. save it, if it supports that sort of thing.
  360. """
  361. params = {"action": "logout"}
  362. self._api_query(params)
  363. self._cookiejar.clear()
  364. self._save_cookiejar()
  365. def _sql_connect(self, **kwargs):
  366. """Attempt to establish a connection with this site's SQL database.
  367. oursql.connect() will be called with self._sql_data as its kwargs,
  368. which is usually config.wiki["sites"][self.name()]["sql"]. Any kwargs
  369. given to this function will be passed to connect() and will have
  370. precedence over the config file.
  371. Will raise SQLError() if the module "oursql" is not available. oursql
  372. may raise its own exceptions (e.g. oursql.InterfaceError) if it cannot
  373. establish a connection.
  374. """
  375. if not oursql:
  376. e = "Module 'oursql' is required for SQL queries."
  377. raise SQLError(e)
  378. args = self._sql_data
  379. for key, value in kwargs.iteritems():
  380. args[key] = value
  381. if "read_default_file" not in args and "user" not in args and "passwd" not in args:
  382. args["read_default_file"] = expanduser("~/.my.cnf")
  383. self._sql_conn = oursql.connect(**args)
  384. def name(self):
  385. """Returns the Site's name (or "wikiid" in the API), like "enwiki"."""
  386. return self._name
  387. def project(self):
  388. """Returns the Site's project name in lowercase, like "wikipedia"."""
  389. return self._project
  390. def lang(self):
  391. """Returns the Site's language code, like "en" or "es"."""
  392. return self._lang
  393. def domain(self):
  394. """Returns the Site's web domain, like "en.wikipedia.org"."""
  395. return urlparse(self._base_url).netloc
  396. def api_query(self, **kwargs):
  397. """Do an API query with `kwargs` as the parameters.
  398. See _api_query()'s documentation for details.
  399. """
  400. return self._api_query(kwargs)
  401. def sql_query(self, query, params=(), plain_query=False, dict_cursor=False,
  402. cursor_class=None, show_table=False):
  403. """Do an SQL query and yield its results.
  404. If `plain_query` is True, we will force an unparameterized query.
  405. Specifying both params and plain_query will cause an error.
  406. If `dict_cursor` is True, we will use oursql.DictCursor as our cursor,
  407. otherwise the default oursql.Cursor. If `cursor_class` is given, it
  408. will override this option.
  409. If `show_table` is True, the name of the table will be prepended to the
  410. name of the column. This will mainly affect a DictCursor.
  411. Example:
  412. >>> query = "SELECT user_id, user_registration FROM user WHERE user_name = ?"
  413. >>> params = ("The Earwig",)
  414. >>> result1 = site.sql_query(query, params)
  415. >>> result2 = site.sql_query(query, params, dict_cursor=True)
  416. >>> for row in result1: print row
  417. (7418060L, '20080703215134')
  418. >>> for row in result2: print row
  419. {'user_id': 7418060L, 'user_registration': '20080703215134'}
  420. See _sql_connect() for information on how a connection is acquired.
  421. <http://packages.python.org/oursql> has helpful documentation on the
  422. oursql module.
  423. This may raise SQLError() or one of oursql's exceptions
  424. (oursql.ProgrammingError, oursql.InterfaceError, ...) if there were
  425. problems with the query.
  426. """
  427. if not self._sql_conn:
  428. self._sql_connect()
  429. if not cursor_class:
  430. if dict_cursor:
  431. cursor_class = oursql.DictCursor
  432. else:
  433. cursor_class = oursql.Cursor
  434. with self._sql_conn.cursor(cursor_class, show_table=show_table) as cur:
  435. cur.execute(query, params, plain_query)
  436. for result in cur:
  437. yield result
  438. def get_replag(self):
  439. """Return the estimated database replication lag in seconds.
  440. Requires SQL access. This function only makes sense on a replicated
  441. database (e.g. the Wikimedia Toolserver) and on a wiki that receives a
  442. large number of edits (ideally, at least one per second), or the result
  443. may be larger than expected.
  444. """
  445. query = "SELECT NOW() - MAX(rev_timestamp) FROM revision"
  446. result = list(self.sql_query(query))
  447. return result[0][0]
  448. def namespace_id_to_name(self, ns_id, all=False):
  449. """Given a namespace ID, returns associated namespace names.
  450. If all is False (default), we'll return the first name in the list,
  451. which is usually the localized version. Otherwise, we'll return the
  452. entire list, which includes the canonical name.
  453. For example, returns u"Wikipedia" if ns_id=4 and all=False on enwiki;
  454. returns [u"Wikipedia", u"Project"] if ns_id=4 and all=True.
  455. Raises NamespaceNotFoundError if the ID is not found.
  456. """
  457. try:
  458. if all:
  459. return self._namespaces[ns_id]
  460. else:
  461. return self._namespaces[ns_id][0]
  462. except KeyError:
  463. e = "There is no namespace with id {0}.".format(ns_id)
  464. raise NamespaceNotFoundError(e)
  465. def namespace_name_to_id(self, name):
  466. """Given a namespace name, returns the associated ID.
  467. Like namespace_id_to_name(), but reversed. Case is ignored, because
  468. namespaces are assumed to be case-insensitive.
  469. Raises NamespaceNotFoundError if the name is not found.
  470. """
  471. lname = name.lower()
  472. for ns_id, names in self._namespaces.items():
  473. lnames = [n.lower() for n in names] # Be case-insensitive
  474. if lname in lnames:
  475. return ns_id
  476. e = "There is no namespace with name '{0}'.".format(name)
  477. raise NamespaceNotFoundError(e)
  478. def get_page(self, title, follow_redirects=False):
  479. """Returns a Page object for the given title (pagename).
  480. Will return a Category object instead if the given title is in the
  481. category namespace. As Category is a subclass of Page, this should not
  482. cause problems.
  483. Note that this doesn't do any direct checks for existence or
  484. redirect-following - Page's methods provide that.
  485. """
  486. prefixes = self.namespace_id_to_name(NS_CATEGORY, all=True)
  487. prefix = title.split(":", 1)[0]
  488. if prefix != title: # Avoid a page that is simply "Category"
  489. if prefix in prefixes:
  490. return Category(self, title, follow_redirects)
  491. return Page(self, title, follow_redirects)
  492. def get_category(self, catname, follow_redirects=False):
  493. """Returns a Category object for the given category name.
  494. `catname` should be given *without* a namespace prefix. This method is
  495. really just shorthand for get_page("Category:" + catname).
  496. """
  497. prefix = self.namespace_id_to_name(NS_CATEGORY)
  498. pagename = ':'.join((prefix, catname))
  499. return Category(self, pagename, follow_redirects)
  500. def get_user(self, username=None):
  501. """Returns a User object for the given username.
  502. If `username` is left as None, then a User object representing the
  503. currently logged-in (or anonymous!) user is returned.
  504. """
  505. if username is None:
  506. username = self._get_username()
  507. return User(self, username)