1 from __future__ import absolute_import

     2 from __future__ import print_function

     3 

     4 import socket

     5 import re

     6 import sslpsk

     7 import Pyro

     8 from Pyro.core import PyroURI

     9 from Pyro.protocol import _connect_socket,TCPConnection,PYROAdapter

    10 from Pyro.errors import ConnectionDeniedError, ProtocolError

    11 from Pyro.util import Log

    12 

    13 #

    14 # The TLS-PSK adapter that handles SSL connections instead of regular sockets,

    15 # but using Pre Shared Keys instead of Certificates

    16 #

    17 class PYROPSKAdapter(PYROAdapter):

    18     # This is essentialy the same as in Pyro/protocol.py

    19     # only raw_sock wrapping into sock through sslpsk.wrap_socket was added

    20     # Pyro unfortunately doesn't allow cleaner customization

    21     def bindToURI(self,URI):

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

    23             try:

    24                 self.URI=URI

    25 

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

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

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

    29                 sock = sslpsk.wrap_socket(

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

    31                     ciphers="PSK-AES256-GCM-SHA384:PSK-AES256-CBC-SHA")

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

    33 

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

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

    36                 try:

    37                     authChallenge=self.recvAuthChallenge(conn)

    38                 except ProtocolError,x:

    39                     # check if we were denied

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

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

    42                     else:

    43                         raise

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

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

    46                 if msg==self.acceptMSG:

    47                     self.conn=conn

    48                     self.conn.connected=1

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

    50                     if URI.protocol=='PYROLOCPSK':

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

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

    53                     try:

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

    55                     except (KeyError,ValueError):

    56                         raise ConnectionDeniedError('invalid response')

    57             except socket.error:

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

    59                 raise ProtocolError('connection failed')

    60 

    61 _getProtocolAdapter = Pyro.protocol.getProtocolAdapter

    62 def getProtocolAdapter(protocol):

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

    64         return PYROPSKAdapter()

    65     _getProtocolAdapter(protocol)

    66 

    67 Pyro.protocol.getProtocolAdapter = getProtocolAdapter

    68 

    69 _processStringURI = Pyro.core.processStringURI

    70 def processStringURI(URI):

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

    72     if x:

    73         protocol=x.group('protocol')

    74         hostname=x.group('hostname')

    75         port=x.group('port')

    76         if port:

    77             port=int(port)

    78         else:

    79             port=0

    80         name=x.group('name')

    81         return PyroURI(hostname,name,port,protocol)

    82     return _processStringURI(URI)

    83 Pyro.core.processStringURI = processStringURI