wiki.py 4.59 KB
Newer Older
Robin Sonnabend's avatar
Robin Sonnabend committed
1
import requests
2
from json import JSONDecodeError
Robin Sonnabend's avatar
Robin Sonnabend committed
3 4 5 6

import config

HTTP_STATUS_OK = 200
7
HTTP_STATUS_AUTHENTICATE = 401
Robin Sonnabend's avatar
Robin Sonnabend committed
8

Robin Sonnabend's avatar
Robin Sonnabend committed
9

Robin Sonnabend's avatar
Robin Sonnabend committed
10 11 12
class WikiException(Exception):
    pass

Robin Sonnabend's avatar
Robin Sonnabend committed
13

Robin Sonnabend's avatar
Robin Sonnabend committed
14 15 16 17 18 19 20 21 22 23
def _filter_params(params):
    result = {}
    for key, value in sorted(params.items(), key=lambda t: t[0] == "token"):
        if isinstance(value, bool):
            if value:
                result[key] = ""
        else:
            result[key] = value
    return result

Robin Sonnabend's avatar
Robin Sonnabend committed
24

Robin Sonnabend's avatar
Robin Sonnabend committed
25
class WikiClient:
Robin Sonnabend's avatar
Robin Sonnabend committed
26 27 28 29 30 31 32 33 34 35 36 37
    def __init__(self, active=None, endpoint=None, anonymous=None, user=None,
                 password=None, domain=None):
        def _or_default(value, default):
            if value is None:
                return default
            return value
        self.active = _or_default(active, config.WIKI_ACTIVE)
        self.endpoint = _or_default(endpoint, config.WIKI_API_URL)
        self.anonymous = _or_default(anonymous, config.WIKI_ANONYMOUS)
        self.user = _or_default(user, config.WIKI_USER)
        self.password = _or_default(password, config.WIKI_PASSWORD)
        self.domain = _or_default(domain, config.WIKI_DOMAIN)
Robin Sonnabend's avatar
Robin Sonnabend committed
38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55
        self.token = None
        self.cookies = requests.cookies.RequestsCookieJar()

    def __enter__(self):
        if not self.anonymous:
            self.login()
        return self

    def __exit__(self, type, value, traceback):
        if not self.anonymous:
            self.logout()

    def is_logged_in(self):
        return self.token is not None

    def login(self):
        if not self.active:
            return
Robin Sonnabend's avatar
Robin Sonnabend committed
56
        # todo: Change this to the new MediaWiki tokens api
Robin Sonnabend's avatar
Robin Sonnabend committed
57 58 59 60
        token_answer = self.do_action("login", method="post", lgname=self.user)
        if "login" not in token_answer or "token" not in token_answer["login"]:
            raise WikiException("No token in login answer.")
        lgtoken = token_answer["login"]["token"]
Robin Sonnabend's avatar
Robin Sonnabend committed
61 62 63
        login_answer = self.do_action(
            "login", method="post", lgname=self.user, lgpassword=self.password,
            lgdomain=self.domain, lgtoken=lgtoken)
Robin Sonnabend's avatar
Robin Sonnabend committed
64
        if ("login" not in login_answer
Robin Sonnabend's avatar
Robin Sonnabend committed
65 66
                or "result" not in login_answer["login"]
                or login_answer["login"]["result"] != "Success"):
Robin Sonnabend's avatar
Robin Sonnabend committed
67 68 69 70 71 72 73
            raise WikiException("Login not successful.")

    def logout(self):
        if not self.active:
            return
        self.do_action("logout")

Robin Sonnabend's avatar
Robin Sonnabend committed
74 75
    def edit_page(self, title, content, summary, recreate=True,
                  createonly=False):
Robin Sonnabend's avatar
Robin Sonnabend committed
76 77 78
        if not self.active:
            return
        # todo: port to new api once the wiki is updated
Robin Sonnabend's avatar
Robin Sonnabend committed
79 80
        prop_answer = self.do_action(
            "query", method="get", prop="info", intoken="edit", titles=title)
Robin Sonnabend's avatar
Robin Sonnabend committed
81
        if ("query" not in prop_answer
Robin Sonnabend's avatar
Robin Sonnabend committed
82
                or "pages" not in prop_answer["query"]):
Robin Sonnabend's avatar
Robin Sonnabend committed
83 84 85 86 87 88 89 90 91
            raise WikiException("Can't get token for page {}".format(title))
        pages = prop_answer["query"]["pages"]
        edit_token = None
        for page in pages.values():
            if page["title"] == title:
                edit_token = page["edittoken"]
                break
        if edit_token is None:
            raise WikiException("Can't get token for page {}".format(title))
Robin Sonnabend's avatar
Robin Sonnabend committed
92 93
        self.do_action(
            action="edit", method="post", data={"text": content},
Robin Sonnabend's avatar
Robin Sonnabend committed
94 95 96 97 98 99 100 101 102 103
            token=edit_token, title=title,
            summary=summary, recreate=recreate,
            createonly=createonly, bot=True)

    def do_action(self, action, method="get", data=None, **kwargs):
        if not self.active:
            return
        kwargs["action"] = action
        kwargs["format"] = "json"
        params = _filter_params(kwargs)
Robin Sonnabend's avatar
Robin Sonnabend committed
104

105 106
        def _do_request():
            if method == "get":
Robin Sonnabend's avatar
Robin Sonnabend committed
107 108 109
                return requests.get(
                    self.endpoint, cookies=self.cookies, params=params,
                    auth=requests.auth.HTTPBasicAuth(self.user, self.password))
110
            elif method == "post":
Robin Sonnabend's avatar
Robin Sonnabend committed
111 112 113 114
                return requests.post(
                    self.endpoint, cookies=self.cookies, data=data,
                    params=params, auth=requests.auth.HTTPBasicAuth(
                        self.user, self.password))
115
        req = _do_request()
Robin Sonnabend's avatar
Robin Sonnabend committed
116
        if req.status_code != HTTP_STATUS_OK:
Robin Sonnabend's avatar
Robin Sonnabend committed
117 118 119
            raise WikiException(
                "HTTP status code {} on action {}.".format(
                    req.status_code, action))
120
        self.cookies.update(req.cookies)
121 122 123 124
        try:
            return req.json()
        except JSONDecodeError:
            raise WikiException("Server did not return valid JSON.")
Robin Sonnabend's avatar
Robin Sonnabend committed
125

Robin Sonnabend's avatar
Robin Sonnabend committed
126

Robin Sonnabend's avatar
Robin Sonnabend committed
127 128
def main():
    with WikiClient() as client:
Robin Sonnabend's avatar
Robin Sonnabend committed
129 130 131
        client.edit_page(
            title="Test", content="This is a very long text.",
            summary="API client test")