uotp/tests/test_uotp.py

import unittest
import sys
import os
import pytest
import pyotp

sys.path.append(
    os.path.dirname(
        os.path.dirname(os.path.abspath(__file__))
    )
)
import app
import config


class TestUOTP(unittest.TestCase):
    app_client = None
    admin_name = 'admin'
    admin_seed = None
    admin_realm = 'ungleich-admin'
    auth_realm = 'ungleich-auth'
    etcd_client = config.etcd_client

    @classmethod
    def setUpClass(cls):
        app.app.config['TESTING'] = True
        os.environ['BASE_PREFIX'] = '/test/uotp/'
        os.environ['ADMIN_REALM'] = TestUOTP.admin_realm
        os.environ['AUTH_REALM'] = TestUOTP.auth_realm

        with app.app.test_client() as client:
            with app.app.app_context():
                app.create_admin_if_dont_exists(TestUOTP.etcd_client)

            entry = TestUOTP.etcd_client.get(
                os.path.join(
                    os.environ['BASE_PREFIX'],
                    'admin'
                    ), value_in_json=True
                )

            assert entry is not None

            TestUOTP.app_client = client
            TestUOTP.admin_seed = entry.value['seed']

        return super().setUpClass()

    @classmethod
    def tearDownClass(cls):
        TestUOTP.etcd_client.client.delete_prefix(os.environ['BASE_PREFIX'])
        del os.environ['BASE_PREFIX']
        del os.environ['ADMIN_REALM']
        del os.environ['AUTH_REALM']
        return super().tearDownClass()

    def get_otp_list(self):
        json_data = {
                        'admin_name': self.admin_name,
                        'admin_realm': self.admin_realm,
                        'admin_token': pyotp.TOTP(self.admin_seed).now()
        }
        r = self.app_client.get('/list', json=json_data)
        return r

    def create_otp(self, _name, _realm,
                   _admin_name, _admin_realm, _admin_seed):

        json_data = {
                        "name": _name,
                        "realm": _realm,
                        "admin_name": _admin_name,
                        "admin_realm": _admin_realm,
                        "admin_token": pyotp.TOTP(_admin_seed).now()
                    }
        r = self.app_client.post('/create', json=json_data)
        return r

    def verify_otp(self, _name, _realm, _seed,
                   _auth_name, _auth_realm, _auth_seed):

        json_data = {
                        "name": _name,
                        "realm": _realm,
                        "token": pyotp.TOTP(_seed).now(),
                        "auth_name": _auth_name,
                        "auth_realm": _auth_realm,
                        "auth_token": pyotp.TOTP(_auth_seed).now()
                    }
        r = self.app_client.get('/verify', json=json_data)
        return r

    def test_list(self):
        r = self.get_otp_list()
        self.assertEqual(r.status_code, 200)
        self.assertIn(self.admin_name, r.json)

    def test_create(self):
        _name = 'auth'
        _realm = 'ungleich-auth'

        # Test Successful case i.e Admin creating OTP Account
        r = self.create_otp(_name, _realm,
                            self.admin_name, self.admin_realm,
                            self.admin_seed)

        self.assertEqual(r.status_code, 200)

        r = self.get_otp_list()
        self.assertIn(_name, r.json)

        # Test Unsuccesful Creation i.e User from non-admin realm
        # tries to create OTP account

        # Get Auth Account
        entry = self.etcd_client.get(
            os.path.join(
                os.environ['BASE_PREFIX'], _name
                ), value_in_json=True
        )
        _adversery_name = 'adversery'
        _adversery_realm = 'ungleich-admin'

        r = self.create_otp(_adversery_name, _adversery_realm,
                            _name, _realm, entry.value['seed'])
        self.assertEqual(r.status_code, 400)

        r = self.get_otp_list()
        self.assertNotIn(
            '{}/{}'.format(_adversery_realm, _adversery_name),
            r.json
        )

    def test_verify(self):
        _auth_name = 'verification-auth'
        _auth_realm = 'ungleich-auth'
        _auth_seed = None

        r = self.create_otp(_auth_name, _auth_realm,
                            self.admin_name, self.admin_realm,
                            self.admin_seed)

        self.assertEqual(r.status_code, 200)
        entry = self.etcd_client.get(
            os.path.join(
                os.environ['BASE_PREFIX'],
                _auth_name
                ),
            value_in_json=True
            )
        self.assertIsNotNone(entry)
        _auth_seed = entry.value['seed']

        # This should work
        r = self.verify_otp(self.admin_name, self.admin_realm,
                            self.admin_seed, _auth_name, _auth_realm,
                            _auth_seed)
        self.assertEqual(r.status_code, 200)

        # This should not work i.e should rerturn 400
        # because the auth_seed is not correct
        r = self.verify_otp(self.admin_name, self.admin_realm,
                            self.admin_seed, _auth_name,
                            _auth_realm, 'meowmeowmeow')
        self.assertEqual(r.status_code, 400)

        # This should not work i.e should rerturn 400
        # because the auth user is not from ungleich-auth realm
        r = self.verify_otp(self.admin_name, self.admin_realm,
                            self.admin_seed, self.admin_name,
                            self.admin_realm, self.admin_seed)
        self.assertEqual(r.status_code, 400)

    def test_delete(self):
        _name = "protem"
        _realm = "protem-realm"
        r = self.create_otp(_name, _realm,
                            self.admin_name,
                            self.admin_realm,
                            self.admin_seed)

        self.assertEqual(r.status_code, 200)

        json_data = {
                        "name": _name,
                        "realm": _realm,
                        'admin_name': self.admin_name,
                        'admin_realm': self.admin_realm,
                        'admin_token': pyotp.TOTP(self.admin_seed).now()
        }

        r = self.app_client.post('/delete', json=json_data)
        self.assertEqual(r.status_code, 200)
        self.assertNotIn(_name, self.get_otp_list().get_json())


if __name__ == '__main__':
    unittest.main()