<gh_stars>1-10
import datetime

from defender.utils import unblock_username, reset_failed_attempts
from django.test import TestCase
from django.urls import reverse
from django_otp.oath import totp
from django_otp.plugins.otp_totp.models import TOTPDevice
from django_otp.util import random_hex

from authentication_service.models import CoreUser as User


class BaseTestCase(TestCase):

    @classmethod
    def setUpTestData(cls):
        super(BaseTestCase, cls).setUpTestData()
        cls.wizard_url = reverse("login")

        cls.standard_user = User.objects.create(
            username="standard_user", birth_date=datetime.date(2000, 1, 1)
        )
        cls.standard_user.set_password("<PASSWORD>")
        cls.standard_user.save()

        cls.twofa_user = User.objects.create(
            username="2fa_user", birth_date=datetime.date(2000, 1, 1)
        )
        cls.twofa_user.set_password("<PASSWORD>")
        cls.twofa_user.save()

        cls.totp_device = TOTPDevice.objects.create(
            user=cls.twofa_user,
            name="default",
            confirmed=True,
            key=random_hex().decode()
        )

        cls.super_user = User.objects.create_superuser(
            username="super_user", email="<EMAIL>",
            password="<PASSWORD>", birth_date=datetime.date(2000, 1, 1)
        )
        cls.super_user.save()

    def setUp(self):
        # Make sure none of the test users are blocked before running a test.
        for user in [self.standard_user, self.twofa_user, self.super_user]:
            unblock_username(user.username)
            reset_failed_attempts(username=user.username)

    def get_credential_step(self):
        print("Getting credential step")
        response = self.client.get(
            self.wizard_url,
            follow=True
        )
        return response

    def post_credential_step(self, user, password):
        print("Posting credential step")
        response = self.client.post(
            self.wizard_url,
            {
                "login_view-current_step": "auth",
                "auth-username": user,
                "auth-password": password
            },
            follow=True
        )
        return response

    def post_token_step(self):
        print("Posting token step")
        response = self.client.post(
            self.wizard_url,
            {
                "login_view-current_step": "token",
                "token-otp_token": totp(self.totp_device.bin_key)
            },
            follow=True
        )
        return response

    def post_backup_step(self, token):
        print("Posting backup step")
        response = self.client.post(
            self.wizard_url,
            {
                "login_view-current_step": "backup",
                "backup-otp_token": token
            },
            follow=True
        )
        return response


class StandardTestCase(BaseTestCase):
    """Test case for users without 2FA enabled."""
    def test_it(self):
        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step. We check for a username field in the response
        # since only staff members will actually have access to admin, other
        # users will see another login screen.
        response = self.post_credential_step(self.standard_user, "1234")
        self.assertEqual(response.status_code, 200)
        self.assertRedirects(response, "{}".format(reverse("edit_profile")))


class TwoFATestCase(BaseTestCase):
    """Test case for users with 2FA enabled."""
    def test_it(self):
        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step
        response = self.post_credential_step(self.twofa_user, "1234")
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Token")

        # Post token step
        response = self.post_token_step()
        self.assertEqual(response.status_code, 200)
        self.assertRedirects(response, "{}".format(reverse("edit_profile")))


class BackupCodeTestCase(BaseTestCase):
    """Test case where user uses one of their backup codes."""

    def test_it(self):
        # Create backup tokens for user
        device = self.twofa_user.staticdevice_set.create(name='backup')
        device.token_set.create(token='<KEY>')

        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step
        response = self.post_credential_step(self.twofa_user, "1234")
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "backup token")

        # Should be able to go to backup tokens step in wizard
        response = self.client.post(
            self.wizard_url, {"wizard_goto_step": "backup"}, follow=True
        )
        self.assertContains(response, "Token")

        # Don't accept invalid tokens
        response = self.client.post(
            self.wizard_url,
            {
                "backup-otp_token": "WRONG",
                "login_view-current_step": "backup"
            },
            follow=True
        )
        self.assertEqual(
            response.context_data["wizard"]["form"].errors,
            {"__all__":
                ["Invalid token. Please make sure you have entered it "
                 "correctly."]
             }
        )

        # Post backup step
        response = self.post_backup_step("abcdef123")
        self.assertEqual(response.status_code, 200)
        self.assertRedirects(response, "{}".format(reverse("edit_profile")))


class SuperUserTestCase(BaseTestCase):
    """Test case for superusers. These users will/should be signed in"""
    def test_it(self):
        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step. Since 2FA isn't set up on the superuser
        # account, they should immediately be signed into the admin interface.
        response = self.post_credential_step(self.super_user, "1234")
        self.assertEquals(response.status_code, 200)
        self.assertRedirects(response, reverse("admin:index"))

import datetime

from defender.utils import unblock_username, reset_failed_attempts
from django.test import TestCase
from django.urls import reverse
from django_otp.oath import totp
from django_otp.plugins.otp_totp.models import TOTPDevice
from django_otp.util import random_hex

from authentication_service.models import CoreUser as User


class BaseTestCase(TestCase):

    @classmethod
    def setUpTestData(cls):
        super(BaseTestCase, cls).setUpTestData()
        cls.wizard_url = reverse("login")

        cls.standard_user = User.objects.create(
            username="standard_user", birth_date=datetime.date(2000, 1, 1)
        )
        cls.standard_user.set_password("<PASSWORD>")
        cls.standard_user.save()

        cls.twofa_user = User.objects.create(
            username="2fa_user", birth_date=datetime.date(2000, 1, 1)
        )
        cls.twofa_user.set_password("<PASSWORD>")
        cls.twofa_user.save()

        cls.totp_device = TOTPDevice.objects.create(
            user=cls.twofa_user,
            name="default",
            confirmed=True,
            key=random_hex().decode()
        )

        cls.super_user = User.objects.create_superuser(
            username="super_user", email="<EMAIL>",
            password="<PASSWORD>", birth_date=datetime.date(2000, 1, 1)
        )
        cls.super_user.save()

    def setUp(self):
        # Make sure none of the test users are blocked before running a test.
        for user in [self.standard_user, self.twofa_user, self.super_user]:
            unblock_username(user.username)
            reset_failed_attempts(username=user.username)

    def get_credential_step(self):
        print("Getting credential step")
        response = self.client.get(
            self.wizard_url,
            follow=True
        )
        return response

    def post_credential_step(self, user, password):
        print("Posting credential step")
        response = self.client.post(
            self.wizard_url,
            {
                "login_view-current_step": "auth",
                "auth-username": user,
                "auth-password": password
            },
            follow=True
        )
        return response

    def post_token_step(self):
        print("Posting token step")
        response = self.client.post(
            self.wizard_url,
            {
                "login_view-current_step": "token",
                "token-otp_token": totp(self.totp_device.bin_key)
            },
            follow=True
        )
        return response

    def post_backup_step(self, token):
        print("Posting backup step")
        response = self.client.post(
            self.wizard_url,
            {
                "login_view-current_step": "backup",
                "backup-otp_token": token
            },
            follow=True
        )
        return response


class StandardTestCase(BaseTestCase):
    """Test case for users without 2FA enabled."""
    def test_it(self):
        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step. We check for a username field in the response
        # since only staff members will actually have access to admin, other
        # users will see another login screen.
        response = self.post_credential_step(self.standard_user, "1234")
        self.assertEqual(response.status_code, 200)
        self.assertRedirects(response, "{}".format(reverse("edit_profile")))


class TwoFATestCase(BaseTestCase):
    """Test case for users with 2FA enabled."""
    def test_it(self):
        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step
        response = self.post_credential_step(self.twofa_user, "1234")
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Token")

        # Post token step
        response = self.post_token_step()
        self.assertEqual(response.status_code, 200)
        self.assertRedirects(response, "{}".format(reverse("edit_profile")))


class BackupCodeTestCase(BaseTestCase):
    """Test case where user uses one of their backup codes."""

    def test_it(self):
        # Create backup tokens for user
        device = self.twofa_user.staticdevice_set.create(name='backup')
        device.token_set.create(token='<KEY>')

        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step
        response = self.post_credential_step(self.twofa_user, "1234")
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "backup token")

        # Should be able to go to backup tokens step in wizard
        response = self.client.post(
            self.wizard_url, {"wizard_goto_step": "backup"}, follow=True
        )
        self.assertContains(response, "Token")

        # Don't accept invalid tokens
        response = self.client.post(
            self.wizard_url,
            {
                "backup-otp_token": "WRONG",
                "login_view-current_step": "backup"
            },
            follow=True
        )
        self.assertEqual(
            response.context_data["wizard"]["form"].errors,
            {"__all__":
                ["Invalid token. Please make sure you have entered it "
                 "correctly."]
             }
        )

        # Post backup step
        response = self.post_backup_step("abcdef123")
        self.assertEqual(response.status_code, 200)
        self.assertRedirects(response, "{}".format(reverse("edit_profile")))


class SuperUserTestCase(BaseTestCase):
    """Test case for superusers. These users will/should be signed in"""
    def test_it(self):
        # Get credentials step
        response = self.get_credential_step()
        self.assertEqual(response.status_code, 200)
        self.assertContains(response, "Username")

        # Post credentials step. Since 2FA isn't set up on the superuser
        # account, they should immediately be signed into the admin interface.
        response = self.post_credential_step(self.super_user, "1234")
        self.assertEquals(response.status_code, 200)
        self.assertRedirects(response, reverse("admin:index"))
# Make sure none of the test users are blocked before running a test. Test case for users without 2FA enabled. # Get credentials step # Post credentials step. We check for a username field in the response # since only staff members will actually have access to admin, other # users will see another login screen. Test case for users with 2FA enabled. # Get credentials step # Post credentials step # Post token step Test case where user uses one of their backup codes. # Create backup tokens for user # Get credentials step # Post credentials step # Should be able to go to backup tokens step in wizard # Don't accept invalid tokens # Post backup step Test case for superusers. These users will/should be signed in # Get credentials step # Post credentials step. Since 2FA isn't set up on the superuser # account, they should immediately be signed into the admin interface.