From adc9fccb7d0e984cd780cd1a785911e36a6316b1 Mon Sep 17 00:00:00 2001
From: Kevin Morris <kevr@0cost.org>
Date: Sat, 26 Dec 2020 19:29:19 -0800
Subject: [PATCH] add aurweb.models.ban.Ban ORM mapping

Signed-off-by: Kevin Morris <kevr@0cost.org>
---
 aurweb/models/ban.py | 19 +++++++++++++
 test/test_ban.py     | 63 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)
 create mode 100644 aurweb/models/ban.py
 create mode 100644 test/test_ban.py

diff --git a/aurweb/models/ban.py b/aurweb/models/ban.py
new file mode 100644
index 00000000..be030380
--- /dev/null
+++ b/aurweb/models/ban.py
@@ -0,0 +1,19 @@
+from fastapi import Request
+from sqlalchemy.orm import mapper
+
+from aurweb.schema import Bans
+
+
+class Ban:
+    def __init__(self, **kwargs):
+        self.IPAddress = kwargs.get("IPAddress")
+        self.BanTS = kwargs.get("BanTS")
+
+
+def is_banned(request: Request):
+    from aurweb.db import session
+    ip = request.client.host
+    return session.query(Ban).filter(Ban.IPAddress == ip).first() is not None
+
+
+mapper(Ban, Bans)
diff --git a/test/test_ban.py b/test/test_ban.py
new file mode 100644
index 00000000..de4f5b1b
--- /dev/null
+++ b/test/test_ban.py
@@ -0,0 +1,63 @@
+import warnings
+
+from datetime import datetime, timedelta
+
+import pytest
+
+from sqlalchemy import exc as sa_exc
+
+from aurweb.models.ban import Ban, is_banned
+from aurweb.testing import setup_test_db
+from aurweb.testing.requests import Request
+
+ban = None
+
+request = Request()
+
+
+@pytest.fixture(autouse=True)
+def setup():
+    from aurweb.db import session
+
+    global ban
+
+    setup_test_db("Bans")
+
+    ban = Ban(IPAddress="127.0.0.1",
+              BanTS=datetime.utcnow() + timedelta(seconds=30))
+    session.add(ban)
+    session.commit()
+
+
+def test_ban():
+    assert ban.IPAddress == "127.0.0.1"
+    assert bool(ban.BanTS)
+
+
+def test_invalid_ban():
+    from aurweb.db import session
+
+    with pytest.raises(sa_exc.IntegrityError,
+                       match="NOT NULL constraint failed: Bans.IPAddress"):
+        bad_ban = Ban(BanTS=datetime.utcnow())
+        session.add(bad_ban)
+
+        # We're adding a ban with no primary key; this causes an
+        # SQLAlchemy warnings when committing to the DB.
+        # Ignore them.
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", sa_exc.SAWarning)
+            session.commit()
+
+    # Since we got a transaction failure, we need to rollback.
+    session.rollback()
+
+
+def test_banned():
+    request.client.host = "127.0.0.1"
+    assert is_banned(request)
+
+
+def test_not_banned():
+    request.client.host = "192.168.0.1"
+    assert not is_banned(request)