Python: block network connections for testing purposes? - python

I'm trying to test a package that provides interfaces to a few web services. It has a test suite that is supposed to test most functions without connecting to the internet. However, there are some lingering tests that may attempt to connect to the internet / download data, and I'd like to prevent them from doing so for two reasons: first, to make sure my test suite works if no network connection is available; second, so that I'm not spamming the web services with excess queries.
An obvious solution is to unplug my machine / turn off wireless, but when I'm running tests on a remote machine that obviously doesn't work.
So, my question: Can I block network / port access for a single python process? ("sandbox" it, but just blocking network connections)
(afaict, pysandbox doesn't do this)
EDIT: I'm using py.test so I need a solution that will work with py.test, in case that affects any proposed answers.

Monkey patching socket ought to do it:
import socket
def guard(*args, **kwargs):
raise Exception("I told you not to use the Internet!")
socket.socket = guard
Make sure this runs before any other import.

Update: There is now a pytest plugin that does the same thing as this answer! You can read the answer just to see how things work, but I strongly recommend using the plugin instead of copying-pasting my answer :-) See here: https://github.com/miketheman/pytest-socket
I found Thomas Orozco's answer to be very helpful. Following on keflavich, this is how I integrated into my unit test suite. This works for me with thousands of very different unit test-cases (<100 that need socket though) ... and in and out of doctests.
I posted it here. Including below for convenience. Tested with Python 2.7.5, pytest==2.7.0. (To test for yourself, run py.test --doctest-modules in directory with all 3 files cloned.)
_socket_toggle.py
from __future__ import print_function
import socket
import sys
_module = sys.modules[__name__]
def disable_socket():
""" disable socket.socket to disable the Internet. useful in testing.
.. doctest::
>>> enable_socket()
[!] socket.socket is enabled.
>>> disable_socket()
[!] socket.socket is disabled. Welcome to the desert of the real.
>>> socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Traceback (most recent call last):
...
RuntimeError: I told you not to use the Internet!
>>> enable_socket()
[!] socket.socket is enabled.
>>> enable_socket()
[!] socket.socket is enabled.
>>> disable_socket()
[!] socket.socket is disabled. Welcome to the desert of the real.
>>> socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Traceback (most recent call last):
...
RuntimeError: I told you not to use the Internet!
>>> enable_socket()
[!] socket.socket is enabled.
"""
setattr(_module, '_socket_disabled', True)
def guarded(*args, **kwargs):
if getattr(_module, '_socket_disabled', False):
raise RuntimeError("I told you not to use the Internet!")
else:
# SocketType is a valid public alias of socket.socket,
# we use it here to avoid namespace collisions
return socket.SocketType(*args, **kwargs)
socket.socket = guarded
print(u'[!] socket.socket is disabled. Welcome to the desert of the real.')
def enable_socket():
""" re-enable socket.socket to enable the Internet. useful in testing.
"""
setattr(_module, '_socket_disabled', False)
print(u'[!] socket.socket is enabled.')
conftest.py
# Put this in the conftest.py at the top of your unit tests folder,
# so it's available to all unit tests
import pytest
import _socket_toggle
def pytest_runtest_setup():
""" disable the interet. test-cases can explicitly re-enable """
_socket_toggle.disable_socket()
#pytest.fixture(scope='function')
def enable_socket(request):
""" re-enable socket.socket for duration of this test function """
_socket_toggle.enable_socket()
request.addfinalizer(_socket_toggle.disable_socket)
test_example.py
# Example usage of the py.test fixture in tests
import socket
import pytest
try:
from urllib2 import urlopen
except ImportError:
import urllib3
urlopen = urllib.request.urlopen
def test_socket_disabled_by_default():
# default behavior: socket.socket is unusable
with pytest.raises(RuntimeError):
urlopen(u'https://www.python.org/')
def test_explicitly_enable_socket(enable_socket):
# socket is enabled by pytest fixture from conftest. disabled in finalizer
assert socket.socket(socket.AF_INET, socket.SOCK_STREAM)

Building on the very helpful answers from Thomas Orozco and driftcatcher here is a variant that works with Python's unittest and (after a small change) Django.
All you need to do is inherit your test case class from the enhanced NoSocketTestCase class and any access to the network will be detected and raises the SocketAccessError exception.
And this approach also works with Django. You only need to change the NoSocketTestCase class to inherit from django.test.TestCase instead of unittest.TestCase.
While not strictly answering OP's question I think this might be helpful for anyone who wants to block network access in unit tests.
no_sockets.py
import socket
from unittest import TestCase
class SocketAccessError(Exception):
pass
class NoSocketsTestCase(TestCase):
"""Enhancement of TestCase class that prevents any use of sockets
Will throw the exception SocketAccessError when any code tries to
access network sockets
"""
#classmethod
def setUpClass(cls):
cls.socket_original = socket.socket
socket.socket = cls.guard
return super().setUpClass()
#classmethod
def tearDownClass(cls):
socket.socket = cls.socket_original
return super().tearDownClass()
#staticmethod
def guard(*args, **kwargs):
raise SocketAccessError('Attempted to access network')
test_no_sockets.py
import urllib.request
from .no_sockets import NoSocketsTestCase, SocketAccessError
class TestNoSocketsTestCase(NoSocketsTestCase):
def test_raises_exception_on_attempted_network_access(self):
with self.assertRaises(SocketAccessError):
urllib.request.urlopen('https://www.google.com')

A simple way to put a gag on the requests library:
from unittest import mock
requests_gag = mock.patch(
'requests.Session.request',
mock.Mock(side_effect=RuntimeError(
'Please use the `responses` library to mock HTTP in your tests.'
))
)
with requests_gag:
... # no Internet here

httpretty is a small library that solves this problem.
If you are using Django test runner, write a custom test runner where you disable all 3rd party API calls.
# common/test_runner.py
import httpretty
from django.test.runner import DiscoverRunner
class CustomTestRunner(DiscoverRunner):
def run_tests(self, *args, **kwargs):
with httpretty.enabled(allow_net_connect=False):
return super().run_tests(*args, **kwargs)
add this new test runner to your settings
TEST_RUNNER = "common.test_runner.CustomTestRunner"
And from now on all external API calls have to be mocked or httpretty.errors.UnmockedError will be raised.
If you are using pytest, this fixture should work.
#pytest.fixture
def disable_external_api_calls():
httpretty.enable()
yield
httpretty.disable()

I have a pytest solution. pytest-network lybrary help me on this.
# conftest.py
import pytest
import socket
_original_connect = socket.socket.connect
def patched_connect(*args, **kwargs):
...
# It depends on your testing purpose
# You may want a exception, add here
# If you test unconnectable situations
# it can stay like this
#pytest.fixture
def enable_network():
socket.socket.connect = _original_connect
yield
socket.socket.connect = patched_connect
#pytest.fixture
def disable_network():
socket.socket.connect = patched_connect
yield
socket.socket.connect = _original_connect
# test_internet.py
def test_your_unconnectable_situation(disable_network):
response = request.get('http://stackoverflow.com/')
response.status_code == 400

Related

pytest to insert caplog fixture in test method

I have the following test class for pytest:
class TestConnection(AsyncTestCase):
'''Integration test'''
#gen_test
def test_connecting_to_server(self):
'''Connecting to the TCPserver'''
client = server = None
try:
sock, port = bind_unused_port()
with NullContext():
server = EchoServer()
server.add_socket(sock)
client = IOStream(socket.socket())
#### HERE I WANT TO HAVE THE caplog FIXTURE
with ExpectLog(app_log, '.*decode.*'):
yield client.connect(('localhost', port))
yield client.write(b'hello\n')
# yield client.read_until(b'\n')
yield gen.moment
assert False
finally:
if server is not None:
server.stop()
if client is not None:
client.close()
Within this class apparently ExpectLog is not working so after a day of digging around in pytest's documentation I found that there is this caplog fixture that you can have inserted in you methods in order to access the captured logs. It seems to work if I have a test function to which I add the caplog argument but how do I make the caplog fixture available within the methods of a test class like the one above?
Although you can't pass fixtures as parameters to unittest test methods, you can inject them as instance attributes. Example:
# spam.py
import logging
def eggs():
logging.getLogger().info('bacon')
Test for spam.eggs():
# test_spam.py
import logging
import unittest
import pytest
import spam
class SpamTest(unittest.TestCase):
#pytest.fixture(autouse=True)
def inject_fixtures(self, caplog):
self._caplog = caplog
def test_eggs(self):
with self._caplog.at_level(logging.INFO):
spam.eggs()
assert self._caplog.records[0].message == 'bacon'

Write a unit test with pytest to test a socket

I wrote a small server chat that does very basic things and I would like to write the tests around it. Unfortunately I quite lost regarding. I would need some help to get on the right tracks.
I have a class called Server() and it contains a method called bind_socket(). I would like to write unit test (preferably using pytest) to test the following method:
class Server(Threading.Thread):
""" Server side class
Instanciate a server in a thread.
"""
MAX_WAITING_CONNECTIONS = 10
def __init__(self, host='localhost', port=10000):
""" Constructor of the Server class.
Initialize the instance in a thread.
Args:
host (str): Host to which to connect (default=localhost)
port (int): Port on which to connect (default=10000)
"""
threading.Thread.__init__(self)
self.host = host
self.port = port
self.connections = []
self.running = True
def bind_socket(self, ip=socket.AF_INET, protocol=socket.SOCK_STREAM):
self.server_socket = socket.socket(ip, protocol)
self.server_socket.bind((self.host, self.port))
self.server_socket.listen(self.MAX_WAITING_CONNECTIONS)
self.connections.append(self.server_socket)
I'm wondering what is the best way to write a test for this method as it doesn't return anything. Should I mock it and try to return the number of of call of socket(), bind(), listen() and append() or is it the wrong way to do proceed? I'm quite lost on that, I did many tests either with pytest and unittest, watch conferences and read articles and I still don't have anything working.
Some explanation and/or examples would be greatly appreciated.
Thanks a lot
For each line of bind_socket you should ask yourself the questions:
What if this line didn't exist
(for conditionals... I know you don't have any here) What if this condition was the other way around
Can this line raise exceptions.
You want your tests to cover all these eventualities.
For example, socket.bind can raise an exception if it's already bound, or socket.listen can raise an exception. Do you close the socket afterwards?

Use TLS and Python for authentication

I want to make a little update script for a software that runs on a Raspberry Pi and works like a local server. That should connect to a master server in the web to get software updates and also to verify the license of the software.
For that I set up two python scripts. I want these to connect via a TLS socket. Then the client checks the server certificate and the server checks if it's one of the authorized clients. I found a solution for this using twisted on this page.
Now there is a problem left. I want to know which client (depending on the certificate) is establishing the connection. Is there a way to do this in Python 3 with twisted?
I'm happy with every answer.
In a word: yes, this is quite possible, and all the necessary stuff is
ported to python 3 - I tested all the following under Python 3.4 on my Mac and it seems to
work fine.
The short answer is
"use twisted.internet.ssl.Certificate.peerFromTransport"
but given that a lot of set-up is required to get to the point where that is
possible, I've constructed a fully working example that you should be able to
try out and build upon.
For posterity, you'll first need to generate a few client certificates all
signed by the same CA. You've probably already done this, but so others can
understand the answer and try it out on their own (and so I could test my
answer myself ;-)), they'll need some code like this:
# newcert.py
from twisted.python.filepath import FilePath
from twisted.internet.ssl import PrivateCertificate, KeyPair, DN
def getCAPrivateCert():
privatePath = FilePath(b"ca-private-cert.pem")
if privatePath.exists():
return PrivateCertificate.loadPEM(privatePath.getContent())
else:
caKey = KeyPair.generate(size=4096)
caCert = caKey.selfSignedCert(1, CN="the-authority")
privatePath.setContent(caCert.dumpPEM())
return caCert
def clientCertFor(name):
signingCert = getCAPrivateCert()
clientKey = KeyPair.generate(size=4096)
csr = clientKey.requestObject(DN(CN=name), "sha1")
clientCert = signingCert.signRequestObject(
csr, serialNumber=1, digestAlgorithm="sha1")
return PrivateCertificate.fromCertificateAndKeyPair(clientCert, clientKey)
if __name__ == '__main__':
import sys
name = sys.argv[1]
pem = clientCertFor(name.encode("utf-8")).dumpPEM()
FilePath(name.encode("utf-8") + b".client.private.pem").setContent(pem)
With this program, you can create a few certificates like so:
$ python newcert.py a
$ python newcert.py b
Now you should have a few files you can use:
$ ls -1 *.pem
a.client.private.pem
b.client.private.pem
ca-private-cert.pem
Then you'll want a client which uses one of these certificates, and sends some
data:
# tlsclient.py
from twisted.python.filepath import FilePath
from twisted.internet.endpoints import SSL4ClientEndpoint
from twisted.internet.ssl import (
PrivateCertificate, Certificate, optionsForClientTLS)
from twisted.internet.defer import Deferred, inlineCallbacks
from twisted.internet.task import react
from twisted.internet.protocol import Protocol, Factory
class SendAnyData(Protocol):
def connectionMade(self):
self.deferred = Deferred()
self.transport.write(b"HELLO\r\n")
def connectionLost(self, reason):
self.deferred.callback(None)
#inlineCallbacks
def main(reactor, name):
pem = FilePath(name.encode("utf-8") + b".client.private.pem").getContent()
caPem = FilePath(b"ca-private-cert.pem").getContent()
clientEndpoint = SSL4ClientEndpoint(
reactor, u"localhost", 4321,
optionsForClientTLS(u"the-authority", Certificate.loadPEM(caPem),
PrivateCertificate.loadPEM(pem)),
)
proto = yield clientEndpoint.connect(Factory.forProtocol(SendAnyData))
yield proto.deferred
import sys
react(main, sys.argv[1:])
And finally, a server which can distinguish between them:
# whichclient.py
from twisted.python.filepath import FilePath
from twisted.internet.endpoints import SSL4ServerEndpoint
from twisted.internet.ssl import PrivateCertificate, Certificate
from twisted.internet.defer import Deferred
from twisted.internet.task import react
from twisted.internet.protocol import Protocol, Factory
class ReportWhichClient(Protocol):
def dataReceived(self, data):
peerCertificate = Certificate.peerFromTransport(self.transport)
print(peerCertificate.getSubject().commonName.decode('utf-8'))
self.transport.loseConnection()
def main(reactor):
pemBytes = FilePath(b"ca-private-cert.pem").getContent()
certificateAuthority = Certificate.loadPEM(pemBytes)
myCertificate = PrivateCertificate.loadPEM(pemBytes)
serverEndpoint = SSL4ServerEndpoint(
reactor, 4321, myCertificate.options(certificateAuthority)
)
serverEndpoint.listen(Factory.forProtocol(ReportWhichClient))
return Deferred()
react(main, [])
For simplicity's sake we'll just re-use the CA's own certificate for the
server, but in a more realistic scenario you'd obviously want a more
appropriate certificate.
You can now run whichclient.py in one window, then python tlsclient.py a;
python tlsclient.py b in another window, and see whichclient.py print out
a and then b respectively, identifying the clients by the commonName
field in their certificate's subject.
The one caveat here is that you might initially want to put that call to
Certificate.peerFromTransport into a connectionMade method; that won't
work.
Twisted does not presently have a callback for "TLS handshake complete";
hopefully it will eventually, but until it does, you have to wait until you've
received some authenticated data from the peer to be sure the handshake has
completed. For almost all applications, this is fine, since by the time you
have received instructions to do anything (download updates, in your case) the
peer must already have sent the certificate.

Best way to run remote commands thru ssh in Twisted?

I have a twisted application which now needs to monitor processes running on several boxes. The way I manually do is 'ssh and ps', now I'd like my twisted application to do. I have 2 options.
Use paramiko or leverage the power of twisted.conch
I really want to use twisted.conch but my research led me to believe that its primarily intended to create SSHServers and SSHClients. However my requirement is a simple remoteExecute(some_cmd)
I was able to figure out how to do this using paramiko but I didnt want to stickparamiko in my twisted app before looking at how to do this using twisted.conch
Code snippets using twisted on how to run remote_cmds using ssh would be highly appreciated. Thanks.
Followup - Happily, the ticket I referenced below is now resolved. The simpler API will be included in the next release of Twisted. The original answer is still a valid way to use Conch and may reveal some interesting details about what's going on, but from Twisted 13.1 and on, if you just want to run a command and handle it's I/O, this simpler interface will work.
It takes an unfortunately large amount of code to execute a command on an SSH using the Conch client APIs. Conch makes you deal with a lot of different layers, even if you just want sensible boring default behavior. However, it's certainly possible. Here's some code which I've been meaning to finish and add to Twisted to simplify this case:
import sys, os
from zope.interface import implements
from twisted.python.failure import Failure
from twisted.python.log import err
from twisted.internet.error import ConnectionDone
from twisted.internet.defer import Deferred, succeed, setDebugging
from twisted.internet.interfaces import IStreamClientEndpoint
from twisted.internet.protocol import Factory, Protocol
from twisted.conch.ssh.common import NS
from twisted.conch.ssh.channel import SSHChannel
from twisted.conch.ssh.transport import SSHClientTransport
from twisted.conch.ssh.connection import SSHConnection
from twisted.conch.client.default import SSHUserAuthClient
from twisted.conch.client.options import ConchOptions
# setDebugging(True)
class _CommandTransport(SSHClientTransport):
_secured = False
def verifyHostKey(self, hostKey, fingerprint):
return succeed(True)
def connectionSecure(self):
self._secured = True
command = _CommandConnection(
self.factory.command,
self.factory.commandProtocolFactory,
self.factory.commandConnected)
userauth = SSHUserAuthClient(
os.environ['USER'], ConchOptions(), command)
self.requestService(userauth)
def connectionLost(self, reason):
if not self._secured:
self.factory.commandConnected.errback(reason)
class _CommandConnection(SSHConnection):
def __init__(self, command, protocolFactory, commandConnected):
SSHConnection.__init__(self)
self._command = command
self._protocolFactory = protocolFactory
self._commandConnected = commandConnected
def serviceStarted(self):
channel = _CommandChannel(
self._command, self._protocolFactory, self._commandConnected)
self.openChannel(channel)
class _CommandChannel(SSHChannel):
name = 'session'
def __init__(self, command, protocolFactory, commandConnected):
SSHChannel.__init__(self)
self._command = command
self._protocolFactory = protocolFactory
self._commandConnected = commandConnected
def openFailed(self, reason):
self._commandConnected.errback(reason)
def channelOpen(self, ignored):
self.conn.sendRequest(self, 'exec', NS(self._command))
self._protocol = self._protocolFactory.buildProtocol(None)
self._protocol.makeConnection(self)
def dataReceived(self, bytes):
self._protocol.dataReceived(bytes)
def closed(self):
self._protocol.connectionLost(
Failure(ConnectionDone("ssh channel closed")))
class SSHCommandClientEndpoint(object):
implements(IStreamClientEndpoint)
def __init__(self, command, sshServer):
self._command = command
self._sshServer = sshServer
def connect(self, protocolFactory):
factory = Factory()
factory.protocol = _CommandTransport
factory.command = self._command
factory.commandProtocolFactory = protocolFactory
factory.commandConnected = Deferred()
d = self._sshServer.connect(factory)
d.addErrback(factory.commandConnected.errback)
return factory.commandConnected
class StdoutEcho(Protocol):
def dataReceived(self, bytes):
sys.stdout.write(bytes)
sys.stdout.flush()
def connectionLost(self, reason):
self.factory.finished.callback(None)
def copyToStdout(endpoint):
echoFactory = Factory()
echoFactory.protocol = StdoutEcho
echoFactory.finished = Deferred()
d = endpoint.connect(echoFactory)
d.addErrback(echoFactory.finished.errback)
return echoFactory.finished
def main():
from twisted.python.log import startLogging
from twisted.internet import reactor
from twisted.internet.endpoints import TCP4ClientEndpoint
# startLogging(sys.stdout)
sshServer = TCP4ClientEndpoint(reactor, "localhost", 22)
commandEndpoint = SSHCommandClientEndpoint("/bin/ls", sshServer)
d = copyToStdout(commandEndpoint)
d.addErrback(err, "ssh command / copy to stdout failed")
d.addCallback(lambda ignored: reactor.stop())
reactor.run()
if __name__ == '__main__':
main()
Some things to note about it:
It uses the new endpoint APIs introduced in Twisted 10.1. It's possible to do this directly on reactor.connectTCP, but I did it as an endpoint to make it more useful; endpoints can be swapped easily without the code that actually asks for a connection knowing.
It does no host key verification at all! _CommandTransport.verifyHostKey is where you would implement that. Take a look at twisted/conch/client/default.py for some hints about what kinds of things you might want to do.
It takes $USER to be the remote username, which you may want to be a parameter.
It probably only works with key authentication. If you want to enable password authentication, you probably need to subclass SSHUserAuthClient and override getPassword to do something.
Almost all of the layers of SSH and Conch are visible here:
_CommandTransport is at the bottom, a plain old protocol that implements the SSH transport protocol. It creates a...
_CommandConnection which implements the SSH connection negotiation parts of the protocol. Once that completes, a...
_CommandChannel is used to talk to a newly opened SSH channel. _CommandChannel does the actual exec to launch your command. Once the channel is opened it creates an instance of...
StdoutEcho, or whatever other protocol you supply. This protocol will get the output from the command you execute, and can write to the command's stdin.
See http://twistedmatrix.com/trac/ticket/4698 for progress in Twisted on supporting this with less code.

Twisted daemon unittest

I've started a client/server project at work using Twisted (I'm a newcomer, so not much experience). I probably did setup things the wrong way/order, because now I'm a little stuck with a Daemon server (using twistd --python to launch it).
I'm wondering if I've to re-implement the server as a standard process to use it in my unittest module?
Here's part of the code to launch the server as a Daemon in the server module (you'll probably recognize part of krondo's articles in this):
class TwistedHawkService(service.Service):
def startService(self):
''''''
service.Service.startService(self)
log.msg('TwistedHawkService running ...')
# Configuration
port = 10000
iface = 'localhost'
topService = service.MultiService()
thService = TwistedHawkService()
thService.setServiceParent(topService)
factory = ReceiverFactory(thService)
tcpService = internet.TCPServer(port, factory, interface=iface)
tcpService.setServiceParent(topService)
application = service.Application("TwistedHawkService")
topService.setServiceParent(application)
I tried copy/pasting the configuration part in the setUp method:
from mfxTwistedHawk.client import mfxTHClient
from mfxTwistedHawk.server import mfxTHServer
class RequestTestCase(TestCase):
def setUp(self):
# Configuration
port = 10000
iface = 'localhost'
self.topService = service.MultiService()
thService = mfxTHServer.TwistedHawkService()
thService.setServiceParent(self.topService)
factory = mfxTHServer.ReceiverFactory(thService)
tcpService = internet.TCPServer(port, factory, interface=iface)
tcpService.setServiceParent(self.topService)
application = service.Application("TwistedHawkService")
self.topService.setServiceParent(application)
def test_connection(self):
mfxTHClient.requestMain('someRequest')
... but of course using trial unittest.py doesn't start it a daemon, so my client can't reach it.
Any advice of how to setup things would be appreciated.
Thanks!
Edit:
Managed to make everything works with this and this, but still feel unsure about the whole thing:
def setUp(self):
# Configuration
port = 10000
iface = 'localhost'
service = mfxTHServer.TwistedHawkService()
factory = mfxTHServer.ReceiverFactory(service)
self.server = reactor.listenTCP(port, factory, interface=iface)
Is it ok to have a daemon implementation for production and standard process for unittest?
Is it ok to have a daemon implementation for production and standard process for unittest?
Your unit test isn't for Twisted's daemonization functionality. It's for the custom application/protocol/server/whatever functionality that you implemented. In general, in a unit test, you want to involve as little code as possible. So in general, it's quite okay, and even preferable, to have your unit tests not daemonize. In fact, you probably want to write some unit tests that don't even listen on a real TCP port, but just call methods on your service, factory, and protocol classes.

Categories

Resources