test_haas_proxy.py 1.66 KB
Newer Older
Michal Horejsek's avatar
Michal Horejsek committed
1 2 3 4 5 6 7 8 9 10
# -*- encoding: utf-8 -*-

import json
try:
    from unittest import mock
except ImportError:
    import mock

import pytest

11 12
from haas_proxy.proxy import ProxySSHSession
from haas_proxy.utils import force_text
Michal Horejsek's avatar
Michal Horejsek committed
13 14 15 16 17


@pytest.fixture
def cmd_args():
    return mock.Mock(
Michal Horejsek's avatar
Michal Horejsek committed
18
        device_token=42,
Michal Horejsek's avatar
Michal Horejsek committed
19 20 21 22 23 24 25 26 27 28 29 30 31 32
        honeypot_host='localhost',
        honeypot_port=2222,
    )


@pytest.fixture(params=(
    ('user', 'pass'),
    (b'user', b'pass'),
))
def avatar(request):
    avatar = mock.Mock(
        username=request.param[0],
        password=request.param[1],
    )
Lukas Nemec's avatar
Lukas Nemec committed
33 34
    avatar.conn.transport.transport.getPeer.return_value = mock.Mock(
        host='hacker', port=12345)
Michal Horejsek's avatar
Michal Horejsek committed
35 36 37 38 39 40 41
    return avatar


@pytest.fixture
def proxy_ssh_session(cmd_args, avatar):
    session = ProxySSHSession(avatar)
    session.cmd_args = cmd_args
Lukas Nemec's avatar
Lukas Nemec committed
42
    session.balancer = mock.Mock(host="localhost", port=2222)
Michal Horejsek's avatar
Michal Horejsek committed
43 44 45 46 47 48 49 50 51 52 53 54 55 56
    return session


@pytest.mark.parametrize('value, expected', (
    ('abc', 'abc'),
    (b'abc', 'abc'),
    ('háčkyčárky', 'háčkyčárky'),
    (b'h\xc3\xa1\xc4\x8dky\xc4\x8d\xc3\xa1rky', 'háčkyčárky'),
))
def test_force_text(value, expected):
    assert force_text(value) == expected


def test_honeypot_ssh_arguments(proxy_ssh_session):
Michal Horejsek's avatar
Michal Horejsek committed
57 58 59 60
    assert proxy_ssh_session.honeypot_ssh_arguments[3:] == [
        'ssh',
        '-o', 'UserKnownHostsFile=/dev/null',
        '-o', 'StrictHostKeyChecking=no',
61
        '-o', 'LogLevel=error',
Michal Horejsek's avatar
Michal Horejsek committed
62 63 64
        '-p', '2222',
        'user@localhost',
    ]
Michal Horejsek's avatar
Michal Horejsek committed
65 66 67 68


def test_mangle_password(proxy_ssh_session):
    assert json.loads(proxy_ssh_session.mangled_password) == {
Michal Horejsek's avatar
Michal Horejsek committed
69 70
        'pass': 'pass',
        'device_token': 42,
Michal Horejsek's avatar
Michal Horejsek committed
71 72 73
        'remote': 'hacker',
        'remote_port': 12345,
    }