A Python robot that edits Wikipedia and interacts with people over IRC https://en.wikipedia.org/wiki/User:EarwigBot
選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。

1014 行
41 KiB

  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright (C) 2009-2024 Ben Kurtovic <ben.kurtovic@gmail.com>
  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 http.cookiejar import CookieJar
  23. from json import dumps
  24. from logging import getLogger, NullHandler
  25. from os.path import expanduser
  26. from threading import RLock
  27. from time import sleep, time
  28. from urllib.parse import unquote_plus, urlparse
  29. import requests
  30. from requests_oauthlib import OAuth1
  31. from earwigbot import exceptions, importer
  32. from earwigbot.wiki import constants
  33. from earwigbot.wiki.category import Category
  34. from earwigbot.wiki.page import Page
  35. from earwigbot.wiki.user import User
  36. pymysql = importer.new("pymysql")
  37. __all__ = ["Site"]
  38. class Site:
  39. """
  40. **EarwigBot: Wiki Toolset: Site**
  41. Represents a site, with support for API queries and returning
  42. :py:class:`~earwigbot.wiki.page.Page`,
  43. :py:class:`~earwigbot.wiki.user.User`,
  44. and :py:class:`~earwigbot.wiki.category.Category` objects. The constructor
  45. takes a bunch of arguments and you probably won't need to call it directly,
  46. rather :py:meth:`wiki.get_site() <earwigbot.wiki.sitesdb.SitesDB.get_site>`
  47. for returning :py:class:`Site`
  48. instances, :py:meth:`wiki.add_site()
  49. <earwigbot.wiki.sitesdb.SitesDB.add_site>` for adding new ones to our
  50. database, and :py:meth:`wiki.remove_site()
  51. <earwigbot.wiki.sitesdb.SitesDB.remove_site>` for removing old ones from
  52. our database, should suffice.
  53. *Attributes:*
  54. - :py:attr:`name`: the site's name (or "wikiid"), like ``"enwiki"``
  55. - :py:attr:`project`: the site's project name, like ``"wikipedia"``
  56. - :py:attr:`lang`: the site's language code, like ``"en"``
  57. - :py:attr:`domain`: the site's web domain, like ``"en.wikipedia.org"``
  58. - :py:attr:`url`: the site's URL, like ``"https://en.wikipedia.org"``
  59. *Public methods:*
  60. - :py:meth:`api_query`: does an API query with kwargs as params
  61. - :py:meth:`sql_query`: does an SQL query and yields its results
  62. - :py:meth:`get_maxlag`: returns the internal database lag
  63. - :py:meth:`get_replag`: estimates the external database lag
  64. - :py:meth:`get_token`: gets a token for a specific API action
  65. - :py:meth:`namespace_id_to_name`: returns names associated with an NS id
  66. - :py:meth:`namespace_name_to_id`: returns the ID associated with a NS name
  67. - :py:meth:`get_page`: returns a Page for the given title
  68. - :py:meth:`get_category`: returns a Category for the given title
  69. - :py:meth:`get_user`: returns a User object for the given name
  70. - :py:meth:`delegate`: controls when the API or SQL is used
  71. """
  72. SERVICE_API = 1
  73. SERVICE_SQL = 2
  74. SPECIAL_TOKENS = [
  75. "createaccount",
  76. "deleteglobalaccount",
  77. "login",
  78. "patrol",
  79. "rollback",
  80. "setglobalaccountstatus",
  81. "userrights",
  82. "watch",
  83. ]
  84. def __init__(
  85. self,
  86. name=None,
  87. project=None,
  88. lang=None,
  89. base_url=None,
  90. article_path=None,
  91. script_path=None,
  92. sql=None,
  93. namespaces=None,
  94. login=(None, None),
  95. oauth=None,
  96. cookiejar=None,
  97. user_agent=None,
  98. use_https=True,
  99. assert_edit=None,
  100. maxlag=None,
  101. wait_between_queries=1,
  102. logger=None,
  103. search_config=None,
  104. ):
  105. """Constructor for new Site instances.
  106. This probably isn't necessary to call yourself unless you're building a
  107. Site that's not in your config and you don't want to add it - normally
  108. all you need is wiki.get_site(name), which creates the Site for you
  109. based on your config file and the sites database. We accept a bunch of
  110. kwargs, but the only ones you really "need" are *base_url* and
  111. *script_path*; this is enough to figure out an API url. *login*, a
  112. tuple of (username, password), can be used to log in using the legacy
  113. BotPasswords system; otherwise, a dict of OAuth info should be provided
  114. to *oauth*. *cookiejar* will be used to store cookies, and we'll use a
  115. normal CookieJar if none is given.
  116. First, we'll store the given arguments as attributes, then set up our
  117. requests session. We'll load any of the attributes that weren't given
  118. from the API, and then log in if a username/pass was given and we
  119. aren't already logged in.
  120. """
  121. # Attributes referring to site information, filled in by an API query
  122. # if they are missing (and an API url can be determined):
  123. self._name = name
  124. self._project = project
  125. self._lang = lang
  126. self._base_url = base_url
  127. self._article_path = article_path
  128. self._script_path = script_path
  129. self._namespaces = namespaces
  130. # Attributes used for API queries:
  131. self._use_https = use_https
  132. self._assert_edit = assert_edit
  133. self._maxlag = maxlag
  134. self._wait_between_queries = wait_between_queries
  135. self._max_retries = 6
  136. self._last_query_time = 0
  137. self._tokens = {}
  138. self._api_lock = RLock()
  139. self._api_info_cache = {"maxlag": 0, "lastcheck": 0}
  140. # Attributes used for SQL queries:
  141. if sql:
  142. self._sql_data = sql
  143. else:
  144. self._sql_data = {}
  145. self._sql_conn = None
  146. self._sql_lock = RLock()
  147. self._sql_info_cache = {"replag": 0, "lastcheck": 0, "usable": None}
  148. # Attribute used in copyright violation checks (see CopyrightMixIn):
  149. if search_config:
  150. self._search_config = search_config
  151. else:
  152. self._search_config = {}
  153. # Set up cookiejar and requests session for making API queries:
  154. if cookiejar is not None:
  155. self._cookiejar = cookiejar
  156. else:
  157. self._cookiejar = CookieJar()
  158. self._last_cookiejar_save = None
  159. if not user_agent:
  160. user_agent = constants.USER_AGENT # Set default UA
  161. self._oauth = oauth
  162. self._session = requests.Session()
  163. self._session.cookies = self._cookiejar
  164. self._session.headers["User-Agent"] = user_agent
  165. if oauth:
  166. self._session.auth = OAuth1(
  167. oauth["consumer_token"],
  168. oauth["consumer_secret"],
  169. oauth["access_token"],
  170. oauth["access_secret"],
  171. )
  172. # Set up our internal logger:
  173. if logger:
  174. self._logger = logger
  175. else: # Just set up a null logger to eat up our messages:
  176. self._logger = getLogger("earwigbot.wiki")
  177. self._logger.addHandler(NullHandler())
  178. # Get all of the above attributes that were not specified as arguments:
  179. self._load_attributes()
  180. # If we have a name/pass and the API says we're not logged in, log in:
  181. self._login_info = name, password = login
  182. if not self._oauth and name and password:
  183. logged_in_as = self._get_username_from_cookies()
  184. if not logged_in_as or name.replace("_", " ") != logged_in_as:
  185. self._login(login)
  186. def __repr__(self):
  187. """Return the canonical string representation of the Site."""
  188. res = ", ".join(
  189. (
  190. "Site(name={_name!r}",
  191. "project={_project!r}",
  192. "lang={_lang!r}",
  193. "base_url={_base_url!r}",
  194. "article_path={_article_path!r}",
  195. "script_path={_script_path!r}",
  196. "use_https={_use_https!r}",
  197. "assert_edit={_assert_edit!r}",
  198. "maxlag={_maxlag!r}",
  199. "sql={_sql_data!r}",
  200. "login={0}",
  201. "oauth={1}",
  202. "user_agent={3!r}",
  203. "cookiejar={2})",
  204. )
  205. )
  206. name, password = self._login_info
  207. login = "({0}, {1})".format(repr(name), "hidden" if password else None)
  208. oauth = "hidden" if self._oauth else None
  209. cookies = self._cookiejar.__class__.__name__
  210. if hasattr(self._cookiejar, "filename"):
  211. cookies += "({0!r})".format(getattr(self._cookiejar, "filename"))
  212. else:
  213. cookies += "()"
  214. agent = self.user_agent
  215. return res.format(login, oauth, cookies, agent, **self.__dict__)
  216. def __str__(self):
  217. """Return a nice string representation of the Site."""
  218. res = "<Site {0} ({1}:{2}) at {3}>"
  219. return res.format(self.name, self.project, self.lang, self.domain)
  220. def _unicodeify(self, value, encoding="utf8"):
  221. """Return input as unicode if it's not unicode to begin with."""
  222. if isinstance(value, str):
  223. return value
  224. return str(value, encoding)
  225. def _api_query(
  226. self,
  227. params,
  228. tries=0,
  229. wait=5,
  230. ignore_maxlag=False,
  231. no_assert=False,
  232. ae_retry=True,
  233. ):
  234. """Do an API query with *params* as a dict of parameters.
  235. See the documentation for :py:meth:`api_query` for full implementation
  236. details. *tries*, *wait*, and *ignore_maxlag* are for maxlag;
  237. *no_assert* and *ae_retry* are for AssertEdit.
  238. """
  239. since_last_query = time() - self._last_query_time # Throttling support
  240. if since_last_query < self._wait_between_queries:
  241. wait_time = self._wait_between_queries - since_last_query
  242. log = "Throttled: waiting {0} seconds".format(round(wait_time, 2))
  243. self._logger.debug(log)
  244. sleep(wait_time)
  245. self._last_query_time = time()
  246. url, params = self._build_api_query(params, ignore_maxlag, no_assert)
  247. if "lgpassword" in params:
  248. self._logger.debug("{0} -> <hidden>".format(url))
  249. else:
  250. data = dumps(params)
  251. if len(data) > 1000:
  252. self._logger.debug("{0} -> {1}...".format(url, data[:997]))
  253. else:
  254. self._logger.debug("{0} -> {1}".format(url, data))
  255. try:
  256. response = self._session.post(url, data=params)
  257. response.raise_for_status()
  258. except requests.RequestException as exc:
  259. raise exceptions.APIError("API query failed: {0}".format(exc))
  260. return self._handle_api_result(response, params, tries, wait, ae_retry)
  261. def _request_csrf_token(self, params):
  262. """If possible, add a request for a CSRF token to an API query."""
  263. if params.get("action") == "query":
  264. if params.get("meta"):
  265. if "tokens" not in params["meta"].split("|"):
  266. params["meta"] += "|tokens"
  267. else:
  268. params["meta"] = "tokens"
  269. if params.get("type"):
  270. if "csrf" not in params["type"].split("|"):
  271. params["type"] += "|csrf"
  272. def _build_api_query(self, params, ignore_maxlag, no_assert):
  273. """Given API query params, return the URL to query and POST data."""
  274. if not self._base_url or self._script_path is None:
  275. e = "Tried to do an API query, but no API URL is known."
  276. raise exceptions.APIError(e)
  277. url = self.url + self._script_path + "/api.php"
  278. params["format"] = "json" # This is the only format we understand
  279. if self._assert_edit and not no_assert:
  280. # If requested, ensure that we're logged in
  281. params["assert"] = self._assert_edit
  282. if self._maxlag and not ignore_maxlag:
  283. # If requested, don't overload the servers:
  284. params["maxlag"] = self._maxlag
  285. if "csrf" not in self._tokens:
  286. # If we don't have a CSRF token, try to fetch one:
  287. self._request_csrf_token(params)
  288. return url, params
  289. def _handle_api_result(self, response, params, tries, wait, ae_retry):
  290. """Given an API query response, attempt to return useful data."""
  291. try:
  292. res = response.json()
  293. except ValueError:
  294. e = "API query failed: JSON could not be decoded."
  295. raise exceptions.APIError(e)
  296. if "warnings" in res:
  297. for name, value in res["warnings"].items():
  298. try:
  299. warning = value["warnings"]
  300. except KeyError:
  301. try:
  302. warning = value["*"]
  303. except KeyError:
  304. warning = value
  305. self._logger.warning("API warning: %s: %s", name, warning)
  306. if self._should_save_cookiejar():
  307. self._save_cookiejar()
  308. try:
  309. code = res["error"]["code"]
  310. info = res["error"]["info"]
  311. except (TypeError, KeyError): # If there's no error code/info, return
  312. if "query" in res and "tokens" in res["query"]:
  313. for name, token in res["query"]["tokens"].items():
  314. self._tokens[name.split("token")[0]] = token
  315. return res
  316. if code == "maxlag": # We've been throttled by the server
  317. if tries >= self._max_retries:
  318. e = "Maximum number of retries reached ({0})."
  319. raise exceptions.APIError(e.format(self._max_retries))
  320. tries += 1
  321. msg = 'Server says "{0}"; retrying in {1} seconds ({2}/{3})'
  322. self._logger.info(msg.format(info, wait, tries, self._max_retries))
  323. sleep(wait)
  324. return self._api_query(params, tries, wait * 2, ae_retry=ae_retry)
  325. elif code in ["assertuserfailed", "assertbotfailed"]: # AssertEdit
  326. if ae_retry and all(self._login_info) and not self._oauth:
  327. # Try to log in if we got logged out:
  328. self._login(self._login_info)
  329. if "token" in params: # Fetch a new one; this is invalid now
  330. params["token"] = self.get_token(params["action"])
  331. return self._api_query(params, tries, wait, ae_retry=False)
  332. if not all(self._login_info) and not self._oauth:
  333. e = "Assertion failed, and no login info was provided."
  334. elif code == "assertbotfailed":
  335. e = "Bot assertion failed: we don't have a bot flag!"
  336. else:
  337. e = "User assertion failed due to an unknown issue. Cookie or OAuth problem?"
  338. raise exceptions.PermissionsError("AssertEdit: " + e)
  339. else: # Some unknown error occurred
  340. e = 'API query failed: got error "{0}"; server says: "{1}".'
  341. error = exceptions.APIError(e.format(code, info))
  342. error.code, error.info = code, info
  343. raise error
  344. def _load_attributes(self, force=False):
  345. """Load data about our Site from the API.
  346. This function is called by __init__() when one of the site attributes
  347. was not given as a keyword argument. We'll do an API query to get the
  348. missing data, but only if there actually *is* missing data.
  349. Additionally, you can call this with *force* set to True to forcibly
  350. reload all attributes.
  351. """
  352. # All attributes to be loaded, except _namespaces, which is a special
  353. # case because it requires additional params in the API query:
  354. attrs = [
  355. self._name,
  356. self._project,
  357. self._lang,
  358. self._base_url,
  359. self._article_path,
  360. self._script_path,
  361. ]
  362. params = {"action": "query", "meta": "siteinfo", "siprop": "general"}
  363. if not self._namespaces or force:
  364. params["siprop"] += "|namespaces|namespacealiases"
  365. with self._api_lock:
  366. result = self._api_query(params, no_assert=True)
  367. self._load_namespaces(result)
  368. elif all(attrs): # Everything is already specified and we're not told
  369. return # to force a reload, so do nothing
  370. else: # We're only loading attributes other than _namespaces
  371. with self._api_lock:
  372. result = self._api_query(params, no_assert=True)
  373. res = result["query"]["general"]
  374. self._name = res["wikiid"]
  375. self._project = res["sitename"].lower()
  376. self._lang = res["lang"]
  377. self._base_url = res["server"]
  378. self._article_path = res["articlepath"]
  379. self._script_path = res["scriptpath"]
  380. def _load_namespaces(self, result):
  381. """Fill self._namespaces with a dict of namespace IDs and names.
  382. Called by _load_attributes() with API data as *result* when
  383. self._namespaces was not given as an kwarg to __init__().
  384. """
  385. self._namespaces = {}
  386. for namespace in result["query"]["namespaces"].values():
  387. ns_id = namespace["id"]
  388. name = namespace["*"]
  389. try:
  390. canonical = namespace["canonical"]
  391. except KeyError:
  392. self._namespaces[ns_id] = [name]
  393. else:
  394. if name != canonical:
  395. self._namespaces[ns_id] = [name, canonical]
  396. else:
  397. self._namespaces[ns_id] = [name]
  398. for namespace in result["query"]["namespacealiases"]:
  399. ns_id = namespace["id"]
  400. alias = namespace["*"]
  401. self._namespaces[ns_id].append(alias)
  402. def _get_cookie(self, name, domain):
  403. """Return the named cookie unless it is expired or doesn't exist."""
  404. for cookie in self._cookiejar:
  405. if cookie.name == name and cookie.domain == domain:
  406. if cookie.is_expired():
  407. break
  408. return cookie
  409. def _get_username_from_cookies(self):
  410. """Try to return our username based solely on cookies.
  411. First, we'll look for a cookie named self._name + "Token", like
  412. "enwikiToken". If it exists and isn't expired, we'll assume it's valid
  413. and try to return the value of the cookie self._name + "UserName" (like
  414. "enwikiUserName"). This should work fine on wikis without single-user
  415. login.
  416. If `enwikiToken` doesn't exist, we'll try to find a cookie named
  417. `centralauth_Token`. If this exists and is not expired, we'll try to
  418. return the value of `centralauth_User`.
  419. If we didn't get any matches, we'll return None. Our goal here isn't to
  420. return the most likely username, or what we *want* our username to be
  421. (for that, we'd do self._login_info[0]), but rather to get our current
  422. username without an unnecessary ?action=query&meta=userinfo API query.
  423. """
  424. name = "".join((self._name, "Token"))
  425. cookie = self._get_cookie(name, self.domain)
  426. if cookie:
  427. name = "".join((self._name, "UserName"))
  428. user_name = self._get_cookie(name, self.domain)
  429. if user_name:
  430. return unquote_plus(user_name.value)
  431. for cookie in self._cookiejar:
  432. if cookie.name != "centralauth_Token" or cookie.is_expired():
  433. continue
  434. base = cookie.domain
  435. if base.startswith(".") and not cookie.domain_initial_dot:
  436. base = base[1:]
  437. if self.domain.endswith(base):
  438. user_name = self._get_cookie("centralauth_User", cookie.domain)
  439. if user_name:
  440. return unquote_plus(user_name.value)
  441. def _get_username_from_api(self):
  442. """Do a simple API query to get our username and return it.
  443. This is a reliable way to make sure we are actually logged in, because
  444. it doesn't deal with annoying cookie logic, but it results in an API
  445. query that is unnecessary in some cases.
  446. Called by _get_username() (in turn called by get_user() with no
  447. username argument) when cookie lookup fails, probably indicating that
  448. we are logged out.
  449. """
  450. result = self.api_query(action="query", meta="userinfo")
  451. return result["query"]["userinfo"]["name"]
  452. def _get_username(self):
  453. """Return the name of the current user, whether logged in or not.
  454. First, we'll try to deduce it solely from cookies, to avoid an
  455. unnecessary API query. For the cookie-detection method, see
  456. _get_username_from_cookies()'s docs.
  457. If our username isn't in cookies, then we're either using OAuth or
  458. we're probably not logged in, or something fishy is going on (like
  459. forced logout). If we're using OAuth and a username was configured,
  460. assume it is accurate and use it. Otherwise, do a single API query for
  461. our username (or IP address) and return that.
  462. """
  463. name = self._get_username_from_cookies()
  464. if name:
  465. return name
  466. if self._oauth and self._login_info[0]:
  467. return self._login_info[0]
  468. return self._get_username_from_api()
  469. def _should_save_cookiejar(self):
  470. """Return a bool indicating whether we should save the cookiejar.
  471. This is True if we haven't saved the cookiejar yet this session, or if
  472. our last save was over a day ago.
  473. """
  474. max_staleness = 60 * 60 * 24 # 1 day
  475. if not self._last_cookiejar_save:
  476. return True
  477. return time() - self._last_cookiejar_save > max_staleness
  478. def _save_cookiejar(self):
  479. """Try to save our cookiejar after doing a (normal) login or logout.
  480. Calls the standard .save() method with no filename. Don't fret if our
  481. cookiejar doesn't support saving (CookieJar raises AttributeError,
  482. FileCookieJar raises NotImplementedError) or no default filename was
  483. given (LWPCookieJar and MozillaCookieJar raise ValueError).
  484. """
  485. if hasattr(self._cookiejar, "save"):
  486. try:
  487. getattr(self._cookiejar, "save")()
  488. except (NotImplementedError, ValueError):
  489. pass
  490. self._last_cookiejar_save = time()
  491. def _login(self, login):
  492. """Safely login through the API.
  493. Normally, this is called by __init__() if a username and password have
  494. been provided and no valid login cookies were found. The only other
  495. time it needs to be called is when those cookies expire, which is done
  496. automatically by api_query() if a query fails.
  497. *login* is a (username, password) tuple.
  498. Raises LoginError on login errors (duh), like bad passwords and
  499. nonexistent usernames.
  500. """
  501. self._tokens.clear()
  502. name, password = login
  503. params = {"action": "query", "meta": "tokens", "type": "login"}
  504. with self._api_lock:
  505. result = self._api_query(params, no_assert=True)
  506. try:
  507. token = result["query"]["tokens"]["logintoken"]
  508. except KeyError:
  509. raise exceptions.LoginError("Couldn't get login token")
  510. params = {
  511. "action": "login",
  512. "lgname": name,
  513. "lgpassword": password,
  514. "lgtoken": token,
  515. }
  516. with self._api_lock:
  517. result = self._api_query(params, no_assert=True)
  518. res = result["login"]["result"]
  519. if res == "Success":
  520. self._tokens.clear()
  521. self._save_cookiejar()
  522. return
  523. if res == "Illegal":
  524. e = "The provided username is illegal."
  525. elif res == "NotExists":
  526. e = "The provided username does not exist."
  527. elif res == "EmptyPass":
  528. e = "No password was given."
  529. elif res == "WrongPass" or res == "WrongPluginPass":
  530. e = "The given password is incorrect."
  531. else:
  532. e = "Couldn't login; server says '{0}'.".format(res)
  533. raise exceptions.LoginError(e)
  534. def _logout(self):
  535. """Safely logout through the API.
  536. We'll do a simple API request (api.php?action=logout), clear our
  537. cookiejar (which probably contains now-invalidated cookies) and try to
  538. save it, if it supports that sort of thing.
  539. """
  540. self.api_query(action="logout")
  541. self._cookiejar.clear()
  542. self._save_cookiejar()
  543. def _sql_connect(self, **kwargs):
  544. """Attempt to establish a connection with this site's SQL database.
  545. pymysql.connect() will be called with self._sql_data as its kwargs.
  546. Any kwargs given to this function will be passed to connect() and will
  547. have precedence over the config file.
  548. Will raise SQLError() if the module "pymysql" is not available. pymysql
  549. may raise its own exceptions (e.g. pymysql.InterfaceError) if it cannot
  550. establish a connection.
  551. """
  552. args = self._sql_data
  553. for key, value in kwargs.items():
  554. args[key] = value
  555. if (
  556. "read_default_file" not in args
  557. and "user" not in args
  558. and "passwd" not in args
  559. ):
  560. args["read_default_file"] = expanduser("~/.my.cnf")
  561. elif "read_default_file" in args:
  562. args["read_default_file"] = expanduser(args["read_default_file"])
  563. if "autoping" not in args:
  564. args["autoping"] = True
  565. if "autoreconnect" not in args:
  566. args["autoreconnect"] = True
  567. try:
  568. self._sql_conn = pymysql.connect(**args)
  569. except ImportError:
  570. e = "SQL querying requires the 'pymysql' package: https://pymysql.readthedocs.io/"
  571. raise exceptions.SQLError(e)
  572. def _get_service_order(self):
  573. """Return a preferred order for using services (e.g. the API and SQL).
  574. A list is returned, starting with the most preferred service first and
  575. ending with the least preferred one. Currently, there are only two
  576. services. SERVICE_API will always be included since the API is expected
  577. to be always usable. In normal circumstances, self.SERVICE_SQL will be
  578. first (with the API second), since using SQL directly is easier on the
  579. servers than making web queries with the API. self.SERVICE_SQL will be
  580. second if replag is greater than three minutes (a cached value updated
  581. every two minutes at most), *unless* API lag is also very high.
  582. self.SERVICE_SQL will not be included in the list if we cannot form a
  583. proper SQL connection.
  584. """
  585. now = time()
  586. if now - self._sql_info_cache["lastcheck"] > 120:
  587. self._sql_info_cache["lastcheck"] = now
  588. try:
  589. try:
  590. self._sql_info_cache["replag"] = sqllag = self.get_replag()
  591. except pymysql.Error as exc:
  592. raise exceptions.SQLError(str(exc))
  593. except (exceptions.SQLError, ImportError):
  594. self._sql_info_cache["usable"] = False
  595. return [self.SERVICE_API]
  596. self._sql_info_cache["usable"] = True
  597. else:
  598. if not self._sql_info_cache["usable"]:
  599. return [self.SERVICE_API]
  600. sqllag = self._sql_info_cache["replag"]
  601. if sqllag > 300:
  602. if not self._maxlag:
  603. return [self.SERVICE_API, self.SERVICE_SQL]
  604. if now - self._api_info_cache["lastcheck"] > 300:
  605. self._api_info_cache["lastcheck"] = now
  606. try:
  607. self._api_info_cache["maxlag"] = apilag = self.get_maxlag()
  608. except exceptions.APIError:
  609. self._api_info_cache["maxlag"] = apilag = 0
  610. else:
  611. apilag = self._api_info_cache["maxlag"]
  612. if apilag > self._maxlag:
  613. return [self.SERVICE_SQL, self.SERVICE_API]
  614. return [self.SERVICE_API, self.SERVICE_SQL]
  615. return [self.SERVICE_SQL, self.SERVICE_API]
  616. @property
  617. def name(self):
  618. """The Site's name (or "wikiid" in the API), like ``"enwiki"``."""
  619. return self._name
  620. @property
  621. def project(self):
  622. """The Site's project name in lowercase, like ``"wikipedia"``."""
  623. return self._project
  624. @property
  625. def lang(self):
  626. """The Site's language code, like ``"en"`` or ``"es"``."""
  627. return self._lang
  628. @property
  629. def domain(self):
  630. """The Site's web domain, like ``"en.wikipedia.org"``."""
  631. return urlparse(self._base_url).netloc
  632. @property
  633. def url(self):
  634. """The Site's full base URL, like ``"https://en.wikipedia.org"``."""
  635. url = self._base_url
  636. if url.startswith("//"): # Protocol-relative URLs from 1.18
  637. if self._use_https:
  638. url = "https:" + url
  639. else:
  640. url = "http:" + url
  641. return url
  642. @property
  643. def user_agent(self):
  644. """The User-Agent header sent to the API by the requests session."""
  645. return self._session.headers["User-Agent"]
  646. def api_query(self, **kwargs):
  647. """Do an API query with `kwargs` as the parameters.
  648. This will first attempt to construct an API url from
  649. :py:attr:`self._base_url` and :py:attr:`self._script_path`. We need
  650. both of these, or else we'll raise
  651. :py:exc:`~earwigbot.exceptions.APIError`. If
  652. :py:attr:`self._base_url` is protocol-relative (introduced in MediaWiki
  653. 1.18), we'll choose HTTPS only if :py:attr:`self._user_https` is
  654. ``True``, otherwise HTTP.
  655. We'll encode the given params, adding ``format=json`` along the way, as
  656. well as ``&assert=`` and ``&maxlag=`` based on
  657. :py:attr:`self._assert_edit` and :py:attr:`_maxlag` respectively.
  658. Additionally, we'll sleep a bit if the last query was made fewer than
  659. :py:attr:`self._wait_between_queries` seconds ago. The request is made
  660. through :py:attr:`self._session`, which has cookie support
  661. (:py:attr:`self._cookiejar`) and a ``User-Agent``
  662. (:py:const:`earwigbot.wiki.constants.USER_AGENT`).
  663. Assuming everything went well, we'll gunzip the data (if compressed),
  664. load it as a JSON object, and return it.
  665. If our request failed for some reason, we'll raise
  666. :py:exc:`~earwigbot.exceptions.APIError` with details. If that
  667. reason was due to maxlag, we'll sleep for a bit and then repeat the
  668. query until we exceed :py:attr:`self._max_retries`.
  669. There is helpful MediaWiki API documentation at `MediaWiki.org
  670. <https://www.mediawiki.org/wiki/API>`_.
  671. """
  672. with self._api_lock:
  673. return self._api_query(kwargs)
  674. def sql_query(
  675. self,
  676. query,
  677. params=(),
  678. plain_query=False,
  679. dict_cursor=False,
  680. cursor_class=None,
  681. buffsize=1024,
  682. ):
  683. """Do an SQL query and yield its results.
  684. If *plain_query* is ``True``, we will force an unparameterized query.
  685. Specifying both *params* and *plain_query* will cause an error. If
  686. *dict_cursor* is ``True``, we will use
  687. :py:class:`pymysql.cursors.DictCursor` as our cursor, otherwise the
  688. default :py:class:`pymysql.cursors.Cursor`. If *cursor_class* is given,
  689. it will override this option.
  690. *buffsize* is the size of each memory-buffered group of results, to
  691. reduce the number of conversations with the database; it is passed to
  692. :py:meth:`cursor.fetchmany() <pymysql.cursors.Cursor.fetchmany>`. If
  693. set to ``0```, all results will be buffered in memory at once (this
  694. uses :py:meth:`fetchall() <pymysql.cursors.Cursor.fetchall>`). If set
  695. to ``1``, it is equivalent to using
  696. :py:meth:`fetchone() <pymysql.cursors.Cursor.fetchone>`.
  697. Example usage::
  698. >>> query = "SELECT user_id, user_registration FROM user WHERE user_name = ?"
  699. >>> params = ("The Earwig",)
  700. >>> result1 = site.sql_query(query, params)
  701. >>> result2 = site.sql_query(query, params, dict_cursor=True)
  702. >>> for row in result1: print row
  703. (7418060L, '20080703215134')
  704. >>> for row in result2: print row
  705. {'user_id': 7418060L, 'user_registration': '20080703215134'}
  706. This may raise :py:exc:`~earwigbot.exceptions.SQLError` or one of
  707. pymysql's exceptions (:py:exc:`pymysql.ProgrammingError`,
  708. :py:exc:`pymysql.InterfaceError`, ...) if there were problems with the
  709. query.
  710. See :py:meth:`_sql_connect` for information on how a connection is
  711. acquired. Also relevant is `pymysql's documentation
  712. <https://pymysql.readthedocs.io/>`_ for details on that package.
  713. """
  714. if not cursor_class:
  715. if dict_cursor:
  716. cursor_class = pymysql.cursors.DictCursor
  717. else:
  718. cursor_class = pymysql.cursors.Cursor
  719. klass = cursor_class
  720. with self._sql_lock:
  721. if not self._sql_conn:
  722. self._sql_connect()
  723. with self._sql_conn.cursor(klass) as cur:
  724. cur.execute(query, params, plain_query)
  725. if buffsize:
  726. while True:
  727. group = cur.fetchmany(buffsize)
  728. if not group:
  729. return
  730. for result in group:
  731. yield result
  732. for result in cur.fetchall():
  733. yield result
  734. def get_maxlag(self, showall=False):
  735. """Return the internal database replication lag in seconds.
  736. In a typical setup, this function returns the replication lag *within*
  737. the WMF's cluster, *not* external replication lag affecting the
  738. Toolserver (see :py:meth:`get_replag` for that). This is useful when
  739. combined with the ``maxlag`` API query param (added by config), in
  740. which queries will be halted and retried if the lag is too high,
  741. usually above five seconds.
  742. With *showall*, will return a list of the lag for all servers in the
  743. cluster, not just the one with the highest lag.
  744. """
  745. params = {"action": "query", "meta": "siteinfo", "siprop": "dbrepllag"}
  746. if showall:
  747. params["sishowalldb"] = 1
  748. with self._api_lock:
  749. result = self._api_query(params, ignore_maxlag=True)
  750. if showall:
  751. return [server["lag"] for server in result["query"]["dbrepllag"]]
  752. return result["query"]["dbrepllag"][0]["lag"]
  753. def get_replag(self):
  754. """Return the estimated external database replication lag in seconds.
  755. Requires SQL access. This function only makes sense on a replicated
  756. database (e.g. the Wikimedia Toolserver) and on a wiki that receives a
  757. large number of edits (ideally, at least one per second), or the result
  758. may be larger than expected, since it works by subtracting the current
  759. time from the timestamp of the latest recent changes event.
  760. This may raise :py:exc:`~earwigbot.exceptions.SQLError` or one of
  761. pymysql's exceptions (:py:exc:`pymysql.ProgrammingError`,
  762. :py:exc:`pymysql.InterfaceError`, ...) if there were problems.
  763. """
  764. query = """SELECT UNIX_TIMESTAMP() - UNIX_TIMESTAMP(rc_timestamp) FROM
  765. recentchanges ORDER BY rc_timestamp DESC LIMIT 1"""
  766. result = list(self.sql_query(query))
  767. return int(result[0][0])
  768. def get_token(self, action=None, force=False):
  769. """Return a token for a data-modifying API action.
  770. In general, this will be a CSRF token, unless *action* is in a special
  771. list of non-CSRF tokens. Tokens are cached for the session (until
  772. :meth:`_login` is called again); set *force* to ``True`` to force a new
  773. token to be fetched.
  774. Raises :exc:`.APIError` if there was an API issue.
  775. """
  776. if action not in self.SPECIAL_TOKENS:
  777. action = "csrf"
  778. if action in self._tokens and not force:
  779. return self._tokens[action]
  780. res = self.api_query(action="query", meta="tokens", type=action)
  781. if action not in self._tokens:
  782. err = "Tried to fetch a {0} token, but API returned: {1}"
  783. raise exceptions.APIError(err.format(action, res))
  784. return self._tokens[action]
  785. def namespace_id_to_name(self, ns_id, all=False):
  786. """Given a namespace ID, returns associated namespace names.
  787. If *all* is ``False`` (default), we'll return the first name in the
  788. list, which is usually the localized version. Otherwise, we'll return
  789. the entire list, which includes the canonical name. For example, this
  790. returns ``u"Wikipedia"`` if *ns_id* = ``4`` and *all* is ``False`` on
  791. ``enwiki``; returns ``[u"Wikipedia", u"Project", u"WP"]`` if *ns_id* =
  792. ``4`` and *all* is ``True``.
  793. Raises :py:exc:`~earwigbot.exceptions.NamespaceNotFoundError` if the ID
  794. is not found.
  795. """
  796. try:
  797. if all:
  798. return self._namespaces[ns_id]
  799. else:
  800. return self._namespaces[ns_id][0]
  801. except KeyError:
  802. e = "There is no namespace with id {0}.".format(ns_id)
  803. raise exceptions.NamespaceNotFoundError(e)
  804. def namespace_name_to_id(self, name):
  805. """Given a namespace name, returns the associated ID.
  806. Like :py:meth:`namespace_id_to_name`, but reversed. Case is ignored,
  807. because namespaces are assumed to be case-insensitive.
  808. Raises :py:exc:`~earwigbot.exceptions.NamespaceNotFoundError` if the
  809. name is not found.
  810. """
  811. lname = name.lower()
  812. for ns_id, names in self._namespaces.items():
  813. lnames = [n.lower() for n in names] # Be case-insensitive
  814. if lname in lnames:
  815. return ns_id
  816. e = "There is no namespace with name '{0}'.".format(name)
  817. raise exceptions.NamespaceNotFoundError(e)
  818. def get_page(self, title, follow_redirects=False, pageid=None):
  819. """Return a :py:class:`Page` object for the given title.
  820. *follow_redirects* is passed directly to
  821. :py:class:`~earwigbot.wiki.page.Page`'s constructor. Also, this will
  822. return a :py:class:`~earwigbot.wiki.category.Category` object instead
  823. if the given title is in the category namespace. As
  824. :py:class:`~earwigbot.wiki.category.Category` is a subclass of
  825. :py:class:`~earwigbot.wiki.page.Page`, this should not cause problems.
  826. Note that this doesn't do any direct checks for existence or
  827. redirect-following: :py:class:`~earwigbot.wiki.page.Page`'s methods
  828. provide that.
  829. """
  830. title = self._unicodeify(title)
  831. prefixes = self.namespace_id_to_name(constants.NS_CATEGORY, all=True)
  832. prefix = title.split(":", 1)[0]
  833. if prefix != title: # Avoid a page that is simply "Category"
  834. if prefix in prefixes:
  835. return Category(self, title, follow_redirects, pageid, self._logger)
  836. return Page(self, title, follow_redirects, pageid, self._logger)
  837. def get_category(self, catname, follow_redirects=False, pageid=None):
  838. """Return a :py:class:`Category` object for the given category name.
  839. *catname* should be given *without* a namespace prefix. This method is
  840. really just shorthand for :py:meth:`get_page("Category:" + catname)
  841. <get_page>`.
  842. """
  843. catname = self._unicodeify(catname)
  844. prefix = self.namespace_id_to_name(constants.NS_CATEGORY)
  845. pagename = ":".join((prefix, catname))
  846. return Category(self, pagename, follow_redirects, pageid, self._logger)
  847. def get_user(self, username=None):
  848. """Return a :py:class:`User` object for the given username.
  849. If *username* is left as ``None``, then a
  850. :py:class:`~earwigbot.wiki.user.User` object representing the currently
  851. logged-in (or anonymous!) user is returned.
  852. """
  853. if username:
  854. username = self._unicodeify(username)
  855. else:
  856. username = self._get_username()
  857. return User(self, username, self._logger)
  858. def delegate(self, services, args=None, kwargs=None):
  859. """Delegate a task to either the API or SQL depending on conditions.
  860. *services* should be a dictionary in which the key is the service name
  861. (:py:attr:`self.SERVICE_API <SERVICE_API>` or
  862. :py:attr:`self.SERVICE_SQL <SERVICE_SQL>`), and the value is the
  863. function to call for this service. All functions will be passed the
  864. same arguments the tuple *args* and the dict *kwargs*, which are both
  865. empty by default. The service order is determined by
  866. :py:meth:`_get_service_order`.
  867. Not every service needs an entry in the dictionary. Will raise
  868. :py:exc:`~earwigbot.exceptions.NoServiceError` if an appropriate
  869. service cannot be found.
  870. """
  871. if not args:
  872. args = ()
  873. if not kwargs:
  874. kwargs = {}
  875. order = self._get_service_order()
  876. for srv in order:
  877. if srv in services:
  878. try:
  879. return services[srv](*args, **kwargs)
  880. except exceptions.ServiceError:
  881. continue
  882. raise exceptions.NoServiceError(services)