Added authentication and access control
authorThierry Florac <thierry.florac@onf.fr>
Mon, 05 Mar 2018 12:27:53 +0100
changeset 13 839b61e1531a
parent 12 6c0dceb71fb1
child 14 5d0e80662a82
Added authentication and access control
src/pyams_zmq/process.py
src/pyams_zmq/socket.py
--- a/src/pyams_zmq/process.py	Sat Jan 27 00:40:59 2018 +0100
+++ b/src/pyams_zmq/process.py	Mon Mar 05 12:27:53 2018 +0100
@@ -23,6 +23,7 @@
 from pyams_zmq.interfaces import IZMQProcess
 
 # import packages
+from zmq.auth.thread import ThreadAuthenticator
 from zmq.eventloop import ioloop, zmqstream
 from zope.interface import implementer
 
@@ -35,8 +36,9 @@
     """
 
     socket_type = zmq.REP
+    auth_thread = None
 
-    def __init__(self, bind_addr, handler):
+    def __init__(self, bind_addr, handler, auth=None, clients=None):
         super(ZMQProcess, self).__init__()
 
         self.context = None
@@ -48,10 +50,18 @@
         self.bind_addr = bind_addr
         self.rep_stream = None
         self.handler = handler
+        self.passwords = dict([auth.split(':', 1)]) if auth else None
+        self.clients = clients.split() if clients else None
 
     def setup(self):
         """Creates a :attr:`context` and an event :attr:`loop` for the process."""
-        self.context = zmq.Context()
+        ctx = self.context = zmq.Context()
+        auth = self.auth_thread = ThreadAuthenticator(ctx)
+        auth.start()
+        if self.clients:
+            auth.allow(*self.clients)
+        if self.passwords:
+            auth.configure_plain(domain='*', passwords=self.passwords)
         self.loop = ioloop.IOLoop.instance()
         self.rep_stream, _ = self.stream(self.socket_type, self.bind_addr, bind=True)
         self.initStream()
@@ -71,6 +81,7 @@
         if self.loop is not None:
             self.loop.stop()
             self.loop = None
+        self.auth_thread.stop()
 
     def exit(self, num, frame):
         self.stop()
@@ -108,6 +119,10 @@
         """
         sock = self.context.socket(sock_type)
 
+        # add server authenticator
+        if self.passwords:
+            sock.plain_server = True
+
         # addr may be 'host:port' or ('host', port)
         if isinstance(addr, str):
             addr = addr.split(':')
--- a/src/pyams_zmq/socket.py	Sat Jan 27 00:40:59 2018 +0100
+++ b/src/pyams_zmq/socket.py	Mon Mar 05 12:27:53 2018 +0100
@@ -21,11 +21,16 @@
 import zmq
 
 
-def zmq_socket(address, socket_type=zmq.REQ, linger=0, protocol='tcp'):
-    """Get ØMQ socket"""
+def zmq_socket(address, socket_type=zmq.REQ, linger=0, protocol='tcp', auth=None):
+    """Get ØMQ socket
+
+    auth is given as unicode 'username:password' string and automatically converted to bytes.
+    """
     context = zmq.Context()
     socket = context.socket(socket_type)
     socket.setsockopt(zmq.LINGER, linger)
+    if auth:
+        socket.plain_username, socket.plain_password = auth.encode().split(b':', 1)
     socket.connect('{0}://{1}'.format(protocol, address))
     return socket