1 files changed, 37 insertions, 0 deletions

diff --git a/wikibania/ban/Ban.py b/wikibania/ban/Ban.py
new file mode 100644
index 0000000..e06ca89
--- /dev/null
+++ b/wikibania/ban/Ban.py
@@ -0,0 +1,37 @@
+from datetime import datetime
+
+import pygeoip
+
+
+ISO_TIMESTAMP = "%Y-%m-%dT%H:%M:%SZ"
+
+
+class Ban:
+    def __init__(self, geoip_looker, user=None, timestamp=None, expiry=None, timestamp_format=ISO_TIMESTAMP):
+        self.geoip_looker = geoip_looker
+        self.timestamp_format = timestamp_format
+
+        self.user = user
+        self.timestamp = timestamp
+        self.expiry = expiry
+
+    def items(self):
+        return {
+            "user": self.user,
+            "timestamp": self.timestamp.strftime(ISO_TIMESTAMP),
+            "expiry": self.expiry.strftime(ISO_TIMESTAMP),
+        }
+
+    def hydrate(self, ban_dict):
+        self.user = ban_dict["user"]
+        self.timestamp = datetime.strptime(ban_dict["timestamp"], ISO_TIMESTAMP)
+        self.expiry = datetime.strptime(ban_dict["expiry"], ISO_TIMESTAMP)
+
+    def calc_duration(self):
+        return (self.expiry - self.timestamp).days
+
+    def lookup_country_code(self):
+        try:
+            return self.geoip_looker.country_code_by_addr(self.user).lower()
+        except pygeoip.GeoIPError:
+            return "UNKNOWN"