31

from __future__ import absolute_import

    32

from __future__ import print_function

    33

    34

import socket

    35

import re

    36

import ssl

    37

import sslpsk

    38

import Pyro

    39

from Pyro.core import PyroURI

    40

from Pyro.protocol import _connect_socket, TCPConnection, PYROAdapter

    41

from Pyro.errors import ConnectionDeniedError, ProtocolError

    42

from Pyro.util import Log

    43

    44

    45

class PYROPSKAdapter(PYROAdapter):

    52

    def bindToURI(self, URI):

    53

        with self.lock:   # only 1 thread at a time can bind the URI

    54

            try:

    55

                self.URI = URI

    56

    57

                # This are the statements that differ from Pyro/protocol.py

    58

                raw_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    59

                _connect_socket(raw_sock, URI.address, URI.port, self.timeout)

    60

                sock = sslpsk.wrap_socket(

    61

                    raw_sock, psk=Pyro.config.PYROPSK, server_side=False,

    62

                    ciphers="PSK-AES256-CBC-SHA",  # available in openssl 1.0.2

    63

                    ssl_version=ssl.PROTOCOL_TLSv1)

    64

                # all the rest is the same as in Pyro/protocol.py

    65

    66

                conn = TCPConnection(sock, sock.getpeername())

    67

                # receive the authentication challenge string, and use that to build the actual identification string.

    68

                try:

    69

                    authChallenge = self.recvAuthChallenge(conn)

    70

                except ProtocolError as x:

    71

                    # check if we were denied

    72

                    if hasattr(x, "partialMsg") and x.partialMsg[:len(self.denyMSG)] == self.denyMSG:

    73

                        raise ConnectionDeniedError(Pyro.constants.deniedReasons[int(x.partialMsg[-1])])

    74

                    else:

    75

                        raise

    76

                # reply with our ident token, generated from the ident passphrase and the challenge

    77

                msg = self._sendConnect(sock, self.newConnValidator.createAuthToken(self.ident, authChallenge, conn.addr, self.URI, None))

    78

                if msg == self.acceptMSG:

    79

                    self.conn = conn

    80

                    self.conn.connected = 1

    81

                    Log.msg('PYROAdapter', 'connected to', str(URI))

    82

                    if URI.protocol == 'PYROLOCPSK':

    83

                        self.resolvePYROLOC_URI("PYROPSK")  # updates self.URI

    84

                elif msg[:len(self.denyMSG)] == self.denyMSG:

    85

                    try:

    86

                        raise ConnectionDeniedError(Pyro.constants.deniedReasons[int(msg[-1])])

    87

                    except (KeyError, ValueError):

    88

                        raise ConnectionDeniedError('invalid response')

    89

            except socket.error:

    90

                Log.msg('PYROAdapter', 'connection failed to URI', str(URI))

    91

                raise ProtocolError('connection failed')

    92

    93

    94

_getProtocolAdapter = Pyro.protocol.getProtocolAdapter

    95

    96

    97

def getProtocolAdapter(protocol):

    98

    if protocol in ('PYROPSK', 'PYROLOCPSK'):

    99

        return PYROPSKAdapter()

   100

    return _getProtocolAdapter(protocol)

   101

   102

   103

_processStringURI = Pyro.core.processStringURI

   104

   105

   106

def processStringURI(URI):

   107

    x = re.match(r'(?P<protocol>PYROLOCPSK)://(?P<hostname>[^\s:]+):?(?P<port>\d+)?/(?P<name>\S*)', URI)

   108

    if x:

   109

        protocol = x.group('protocol')

   110

        hostname = x.group('hostname')

   111

        port = x.group('port')

   112

        if port:

   113

            port = int(port)

   114

        else:

   115

            port = 0

   116

        name = x.group('name')

   117

        return PyroURI(hostname, name, port, protocol)

   118

    return _processStringURI(URI)

   119

   120

   121

def setupPSKAdapter():

   122

    """

   123

    Add PyroAdapter to the list of available in

   124

    Pyro adapters and handle new supported protocols

   125

   126

    This function should be called after

   127

    reimport of Pyro module to enable PYROS:// again.

   128

    """

   129

    Pyro.protocol.getProtocolAdapter = getProtocolAdapter

   130

    Pyro.core.processStringURI = processStringURI