1 files changed, 50 insertions, 0 deletions

diff --git a/wikibania/ban/BanDB.py b/wikibania/ban/BanDB.py
new file mode 100644
index 0000000..e83aa3c
--- /dev/null
+++ b/wikibania/ban/BanDB.py
@@ -0,0 +1,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)