summaryrefslogtreecommitdiff
path: root/wikibania/ban/BanDB.py
blob: e83aa3cc8bac704a5182ee2dd7f40184c59fc7e6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import json

from wikibania.ban.Ban import Ban
from wikibania.wapi.WikipediaQuery import BlockQuery


class BanDB:
    def __init__(self, geoip_looker):
        self.geoip_looker = geoip_looker
        self.bans = []

    def list(self):
        return self.bans

    def load(self, ban_list):
        for entry in ban_list:
            ban = Ban(self.geoip_looker)
            ban.hydrate(entry)
            self.bans.append(ban)

    def load_file(self, file_name):
        with open(file_name, "r") as file:
            entries = json.load(file)
            self.load(entries)

    def dump(self):
        return [ban.items() for ban in self.bans]

    def dump_file(self, file_name):
        with open(file_name, "w") as file:
            ban_list = self.dump()
            json.dump(ban_list, file)

    def fetch(self, nb_samples, query_limit=500, continue_token=None):
        fetch = min(nb_samples, query_limit)

        query = BlockQuery(
            properties=["user", "timestamp", "expiry"],
            show=["temp", "ip"],
            limit=fetch,
            continue_token=continue_token,
        )
        results = query.fetch_result()

        entries = results["query"]["blocks"]
        self.load(entries)

        if nb_samples - fetch > 0:
            continue_token = results["query-continue"]["blocks"]["bkcontinue"]
            self.fetch(nb_samples - fetch, query_limit, continue_token)