Add tests for sudo mode

hc/accounts/tests/test_profile.py

@@ -4,6 +4,7 @@ from django.core import mail
 from django.conf import settings
 from django.utils.timezone import now
 from hc.test import BaseTestCase
+from hc.accounts.models import Credential
 from hc.api.models import Check
 
 
@@ -150,3 +151,10 @@ class ProfileTestCase(BaseTestCase):
 
         r = self.client.get("/accounts/profile/")
         self.assertContains(r, "You do not have any projects. Create one!")
+
+    def test_it_shows_security_key(self):
+        Credential.objects.create(user=self.alice, name="Alices Key")
+
+        self.client.login(username="[email protected]", password="password")
+        r = self.client.get("/accounts/profile/")
+        self.assertContains(r, "Alices Key")

hc/accounts/tests/test_sudo_mode.py

@@ -0,0 +1,73 @@
+from django.core import mail
+from django.core.signing import TimestampSigner
+
+from hc.test import BaseTestCase
+from hc.accounts.models import Credential
+from hc.api.models import TokenBucket
+
+
+class SudoModeTestCase(BaseTestCase):
+    def setUp(self):
+        super().setUp()
+
+        self.c = Credential.objects.create(user=self.alice, name="Alices Key")
+        self.url = f"/accounts/two_factor/{self.c.code}/remove/"
+
+    def test_it_sends_code(self):
+        self.client.login(username="[email protected]", password="password")
+
+        r = self.client.get(self.url)
+        self.assertContains(r, "We have sent a confirmation code")
+
+        # A code should have been sent
+        self.assertEqual(len(mail.outbox), 1)
+
+        email = mail.outbox[0]
+        self.assertEqual(email.to[0], "[email protected]")
+        self.assertIn("Confirmation code", email.subject)
+
+    def test_it_accepts_code(self):
+        self.client.login(username="[email protected]", password="password")
+
+        session = self.client.session
+        session["sudo_code"] = TimestampSigner().sign("123456")
+        session.save()
+
+        r = self.client.post(self.url, {"sudo_code": "123456"})
+        self.assertRedirects(r, self.url)
+
+        # sudo mode should now be active
+        self.assertIn("sudo", self.client.session)
+
+    def test_it_rejects_incorrect_code(self):
+        self.client.login(username="[email protected]", password="password")
+
+        session = self.client.session
+        session["sudo_code"] = TimestampSigner().sign("123456")
+        session.save()
+
+        r = self.client.post(self.url, {"sudo_code": "000000"})
+        self.assertContains(r, "Not a valid code.")
+
+        # sudo mode should *not* be active
+        self.assertNotIn("sudo", self.client.session)
+
+    def test_it_passes_through_if_sudo_mode_is_active(self):
+        self.client.login(username="[email protected]", password="password")
+
+        session = self.client.session
+        session["sudo"] = TimestampSigner().sign("active")
+        session.save()
+
+        r = self.client.get(self.url)
+        self.assertContains(r, "Remove Security Key")
+
+    def test_it_uses_rate_limiting(self):
+        self.client.login(username="[email protected]", password="password")
+
+        obj = TokenBucket(value=f"sudo-{self.alice.id}")
+        obj.tokens = 0
+        obj.save()
+
+        r = self.client.get(self.url)
+        self.assertContains(r, "Too Many Requests")