webssh/tests/test_auth_app.py

import json
import shutil
import tempfile

from tornado.options import options
from tornado.testing import AsyncHTTPTestCase

from webssh.handler import login_rate_limiter
from webssh.worker import clients
from webssh.main import make_app, make_handlers
from webssh.settings import get_app_settings
from webssh.utils import to_str


class TestAppAuth(AsyncHTTPTestCase):

    def setUp(self):
        self.tmpdir = tempfile.mkdtemp()
        super(TestAppAuth, self).setUp()

    def tearDown(self):
        super(TestAppAuth, self).tearDown()
        shutil.rmtree(self.tmpdir)

    def get_app(self):
        loop = self.io_loop
        options.auth = True
        options.auth_username = 'admin'
        options.auth_password = ''
        options.auth_password_hash = ''
        options.auth_session_days = 7
        options.data_dir = self.tmpdir
        options.debug = False
        options.xsrf = False
        options.policy = 'warning'
        options.hostfile = ''
        options.syshostfile = ''
        options.tdstream = ''
        options.origin = 'same'
        handlers = make_handlers(loop, options)
        return make_app(handlers, get_app_settings(options))

    def get_cookie_header(self, response):
        cookie = response.headers.get_list('Set-Cookie')[0]
        return {'Cookie': cookie.split(';', 1)[0]}

    def test_auth_setup_and_saved_connections_endpoint(self):
        response = self.fetch('/', follow_redirects=False)
        self.assertEqual(response.code, 302)
        self.assertIn('/login', response.headers['Location'])

        response = self.fetch('/login')
        self.assertIn('创建管理员账号'.encode('utf-8'), response.body)

        body = 'username=admin&password=secret&confirm=secret'
        response = self.fetch('/login', method='POST', body=body,
                              follow_redirects=False)
        self.assertEqual(response.code, 302)
        headers = self.get_cookie_header(response)

        response = self.fetch('/', headers=headers)
        self.assertIn('已保存连接'.encode('utf-8'), response.body)

        response = self.fetch('/connections', headers=headers)
        data = json.loads(to_str(response.body))
        self.assertEqual(data, {'connections': []})

    def test_login_is_rate_limited_after_repeated_failures(self):
        login_rate_limiter.clear()
        self.addCleanup(login_rate_limiter.clear)

        body = 'username=admin&password=secret&confirm=secret'
        response = self.fetch('/login', method='POST', body=body,
                              follow_redirects=False)
        self.assertEqual(response.code, 302)

        wrong = 'username=admin&password=nope'
        for _ in range(login_rate_limiter.max_attempts):
            response = self.fetch('/login', method='POST', body=wrong)
            self.assertIn('用户名或密码错误'.encode('utf-8'), response.body)

        # further attempts are refused without touching the password hash
        response = self.fetch('/login', method='POST', body=wrong)
        self.assertIn('登录尝试过于频繁'.encode('utf-8'), response.body)

        # ... and a correct password is refused too while the block lasts
        good = 'username=admin&password=secret'
        response = self.fetch('/login', method='POST', body=good,
                              follow_redirects=False)
        self.assertEqual(response.code, 200)
        self.assertIn('登录尝试过于频繁'.encode('utf-8'), response.body)

    def test_successful_login_clears_the_failure_counter(self):
        login_rate_limiter.clear()
        self.addCleanup(login_rate_limiter.clear)

        body = 'username=admin&password=secret&confirm=secret'
        self.fetch('/login', method='POST', body=body, follow_redirects=False)

        self.fetch('/login', method='POST', body='username=admin&password=x')
        self.assertNotEqual({}, login_rate_limiter.attempts)

        response = self.fetch('/login', method='POST',
                              body='username=admin&password=secret',
                              follow_redirects=False)
        self.assertEqual(response.code, 302)
        self.assertEqual({}, login_rate_limiter.attempts)


class TestMaxConnPerUser(AsyncHTTPTestCase):

    def setUp(self):
        self.tmpdir = tempfile.mkdtemp()
        clients.clear()
        login_rate_limiter.clear()
        self.addCleanup(clients.clear)
        self.addCleanup(login_rate_limiter.clear)
        super(TestMaxConnPerUser, self).setUp()

    def tearDown(self):
        super(TestMaxConnPerUser, self).tearDown()
        shutil.rmtree(self.tmpdir, ignore_errors=True)

    def get_app(self):
        loop = self.io_loop
        options.auth = True
        options.auth_username = 'admin'
        options.auth_password = 'secret'
        options.auth_password_hash = ''
        options.auth_session_days = 7
        options.data_dir = self.tmpdir
        options.debug = False
        options.xsrf = False
        options.policy = 'warning'
        options.hostfile = ''
        options.syshostfile = ''
        options.tdstream = ''
        options.origin = 'same'
        options.maxconn = 1
        return make_app(make_handlers(loop, options), get_app_settings(options))

    def login(self):
        response = self.fetch('/login', method='POST',
                              body='username=admin&password=secret',
                              follow_redirects=False)
        cookie = response.headers.get_list('Set-Cookie')[0]
        return {'Cookie': cookie.split(';', 1)[0]}

    def test_limit_follows_the_login_not_the_address(self):
        headers = self.login()
        body = 'hostname=127.0.0.1&port=65000&username=robey'

        # one live session for this login already exhausts maxconn=1
        clients['user:admin'] = {'fake': None}
        response = self.fetch('/', method='POST', body=body, headers=headers)
        self.assertIn(b'Too many live connections', response.body)

        # another user shares the same source address but has its own budget,
        # so it gets as far as actually dialing the (absent) ssh server
        clients.clear()
        clients['user:someone-else'] = {'fake': None}
        response = self.fetch('/', method='POST', body=body, headers=headers)
        self.assertNotIn(b'Too many live connections', response.body)
        self.assertIn(b'Unable to connect to', response.body)