sshuttle/tests/client/test_helpers.py
Scott Kuhl 7fc33c0020 Refactor debug, log and Fatal messages.
This commit rewrites the log() function so that it will append a
newline at the end of the message if none is present. It doesn't make
sense to print a log message without a newline since the next log
message (which will write a prefix) expects to be starting at the
beginning of a line.

Although it isn't strictly necessary, this commit also removes any
newlines at the ends of messages. If I missed any, including the
newline at the end of the message will continue to work as it did
before.

Previously, some calls were missing the newline at the end even though
including it was necessary for subsequent messages to appear
correctly.

This code also cleans up some redundant prefixes. The log() method
will prepend the prefix and the different processes should set their
prefix as soon as they start.

Multiline messages are still supported (although the prefix for the
additional lines was changed to match the length of the prefix used
for the first line).
2021-01-01 19:32:48 +11:00

197 lines
6.1 KiB
Python

import io
import socket
from socket import AF_INET, AF_INET6
import errno
from mock import patch, call
import sshuttle.helpers
@patch('sshuttle.helpers.logprefix', new='prefix: ')
@patch('sshuttle.helpers.sys.stdout')
@patch('sshuttle.helpers.sys.stderr')
def test_log(mock_stderr, mock_stdout):
sshuttle.helpers.log("message")
sshuttle.helpers.log("abc")
sshuttle.helpers.log("message 1\n")
sshuttle.helpers.log("message 2\nline2\nline3\n")
sshuttle.helpers.log("message 3\nline2\nline3")
assert mock_stdout.mock_calls == [
call.flush(),
call.flush(),
call.flush(),
call.flush(),
call.flush(),
]
assert mock_stderr.mock_calls == [
call.write('prefix: message\n'),
call.flush(),
call.write('prefix: abc\n'),
call.flush(),
call.write('prefix: message 1\n'),
call.flush(),
call.write('prefix: message 2\n'),
call.write(' line2\n'),
call.write(' line3\n'),
call.flush(),
call.write('prefix: message 3\n'),
call.write(' line2\n'),
call.write(' line3\n'),
call.flush(),
]
@patch('sshuttle.helpers.logprefix', new='prefix: ')
@patch('sshuttle.helpers.verbose', new=1)
@patch('sshuttle.helpers.sys.stdout')
@patch('sshuttle.helpers.sys.stderr')
def test_debug1(mock_stderr, mock_stdout):
sshuttle.helpers.debug1("message")
assert mock_stdout.mock_calls == [
call.flush(),
]
assert mock_stderr.mock_calls == [
call.write('prefix: message\n'),
call.flush(),
]
@patch('sshuttle.helpers.logprefix', new='prefix: ')
@patch('sshuttle.helpers.verbose', new=0)
@patch('sshuttle.helpers.sys.stdout')
@patch('sshuttle.helpers.sys.stderr')
def test_debug1_nop(mock_stderr, mock_stdout):
sshuttle.helpers.debug1("message")
assert mock_stdout.mock_calls == []
assert mock_stderr.mock_calls == []
@patch('sshuttle.helpers.logprefix', new='prefix: ')
@patch('sshuttle.helpers.verbose', new=2)
@patch('sshuttle.helpers.sys.stdout')
@patch('sshuttle.helpers.sys.stderr')
def test_debug2(mock_stderr, mock_stdout):
sshuttle.helpers.debug2("message")
assert mock_stdout.mock_calls == [
call.flush(),
]
assert mock_stderr.mock_calls == [
call.write('prefix: message\n'),
call.flush(),
]
@patch('sshuttle.helpers.logprefix', new='prefix: ')
@patch('sshuttle.helpers.verbose', new=1)
@patch('sshuttle.helpers.sys.stdout')
@patch('sshuttle.helpers.sys.stderr')
def test_debug2_nop(mock_stderr, mock_stdout):
sshuttle.helpers.debug2("message")
assert mock_stdout.mock_calls == []
assert mock_stderr.mock_calls == []
@patch('sshuttle.helpers.logprefix', new='prefix: ')
@patch('sshuttle.helpers.verbose', new=3)
@patch('sshuttle.helpers.sys.stdout')
@patch('sshuttle.helpers.sys.stderr')
def test_debug3(mock_stderr, mock_stdout):
sshuttle.helpers.debug3("message")
assert mock_stdout.mock_calls == [
call.flush(),
]
assert mock_stderr.mock_calls == [
call.write('prefix: message\n'),
call.flush(),
]
@patch('sshuttle.helpers.logprefix', new='prefix: ')
@patch('sshuttle.helpers.verbose', new=2)
@patch('sshuttle.helpers.sys.stdout')
@patch('sshuttle.helpers.sys.stderr')
def test_debug3_nop(mock_stderr, mock_stdout):
sshuttle.helpers.debug3("message")
assert mock_stdout.mock_calls == []
assert mock_stderr.mock_calls == []
@patch('sshuttle.helpers.open', create=True)
def test_resolvconf_nameservers(mock_open):
mock_open.return_value = io.StringIO(u"""
# Generated by NetworkManager
search pri
nameserver 192.168.1.1
nameserver 192.168.2.1
nameserver 192.168.3.1
nameserver 192.168.4.1
nameserver 2404:6800:4004:80c::1
nameserver 2404:6800:4004:80c::2
nameserver 2404:6800:4004:80c::3
nameserver 2404:6800:4004:80c::4
""")
ns = sshuttle.helpers.resolvconf_nameservers(False)
assert ns == [
(AF_INET, u'192.168.1.1'), (AF_INET, u'192.168.2.1'),
(AF_INET, u'192.168.3.1'), (AF_INET, u'192.168.4.1'),
(AF_INET6, u'2404:6800:4004:80c::1'),
(AF_INET6, u'2404:6800:4004:80c::2'),
(AF_INET6, u'2404:6800:4004:80c::3'),
(AF_INET6, u'2404:6800:4004:80c::4')
]
@patch('sshuttle.helpers.open', create=True)
def test_resolvconf_random_nameserver(mock_open):
mock_open.return_value = io.StringIO(u"""
# Generated by NetworkManager
search pri
nameserver 192.168.1.1
nameserver 192.168.2.1
nameserver 192.168.3.1
nameserver 192.168.4.1
nameserver 2404:6800:4004:80c::1
nameserver 2404:6800:4004:80c::2
nameserver 2404:6800:4004:80c::3
nameserver 2404:6800:4004:80c::4
""")
ns = sshuttle.helpers.resolvconf_random_nameserver(False)
assert ns in [
(AF_INET, u'192.168.1.1'), (AF_INET, u'192.168.2.1'),
(AF_INET, u'192.168.3.1'), (AF_INET, u'192.168.4.1'),
(AF_INET6, u'2404:6800:4004:80c::1'),
(AF_INET6, u'2404:6800:4004:80c::2'),
(AF_INET6, u'2404:6800:4004:80c::3'),
(AF_INET6, u'2404:6800:4004:80c::4')
]
@patch('sshuttle.helpers.socket.socket.bind')
def test_islocal(mock_bind):
bind_error = socket.error(errno.EADDRNOTAVAIL)
mock_bind.side_effect = [None, bind_error, None, bind_error]
assert sshuttle.helpers.islocal("127.0.0.1", AF_INET)
assert not sshuttle.helpers.islocal("192.0.2.1", AF_INET)
assert sshuttle.helpers.islocal("::1", AF_INET6)
assert not sshuttle.helpers.islocal("2001:db8::1", AF_INET6)
def test_family_ip_tuple():
assert sshuttle.helpers.family_ip_tuple("127.0.0.1") \
== (AF_INET, "127.0.0.1")
assert sshuttle.helpers.family_ip_tuple("192.168.2.6") \
== (AF_INET, "192.168.2.6")
assert sshuttle.helpers.family_ip_tuple("::1") \
== (AF_INET6, "::1")
assert sshuttle.helpers.family_ip_tuple("2404:6800:4004:80c::1") \
== (AF_INET6, "2404:6800:4004:80c::1")
def test_family_to_string():
assert sshuttle.helpers.family_to_string(AF_INET) == "AF_INET"
assert sshuttle.helpers.family_to_string(AF_INET6) == "AF_INET6"
expected = 'AddressFamily.AF_UNIX'
assert sshuttle.helpers.family_to_string(socket.AF_UNIX) == expected